diff --git a/drivers/hwtracing/coresight/coresight-ost.c b/drivers/hwtracing/coresight/coresight-ost.c
index 07ba40862b4a..943e5c4ea855 100644
--- a/drivers/hwtracing/coresight/coresight-ost.c
+++ b/drivers/hwtracing/coresight/coresight-ost.c
@@ -281,14 +281,13 @@ EXPORT_SYMBOL(stm_ost_packet);
 
 int stm_set_ost_params(struct stm_drvdata *drvdata, size_t bitmap_size)
 {
-	stmdrvdata = drvdata;
-
 	drvdata->chs.bitmap = devm_kzalloc(drvdata->dev, bitmap_size,
 					   GFP_KERNEL);
 	if (!drvdata->chs.bitmap)
 		return -ENOMEM;
 
 	bitmap_fill(drvdata->entities, OST_ENTITY_MAX);
+	stmdrvdata = drvdata;
 
 	return 0;
 }
diff --git a/drivers/hwtracing/coresight/coresight-stm.c b/drivers/hwtracing/coresight/coresight-stm.c
index bc4287ad48b0..bacba944a972 100644
--- a/drivers/hwtracing/coresight/coresight-stm.c
+++ b/drivers/hwtracing/coresight/coresight-stm.c
@@ -870,11 +870,6 @@ static int stm_probe(struct amba_device *adev, const struct amba_id *id)
 				 BYTES_PER_CHANNEL), resource_size(res));
 	}
 	bitmap_size = BITS_TO_LONGS(drvdata->numsp) * sizeof(long);
-	/* Store the driver data pointer for use in exported functions */
-	ret = stm_set_ost_params(drvdata, bitmap_size);
-	if (ret)
-		return ret;
-
 
 	guaranteed = devm_kzalloc(dev, bitmap_size, GFP_KERNEL);
 	if (!guaranteed)
@@ -904,6 +899,11 @@ static int stm_probe(struct amba_device *adev, const struct amba_id *id)
 		goto stm_unregister;
 	}
 
+	/* Store the driver data pointer for use in exported functions */
+	ret = stm_set_ost_params(drvdata, bitmap_size);
+	if (ret)
+		goto stm_unregister;
+
 	pm_runtime_put(&adev->dev);
 
 	dev_info(dev, "%s initialized with master %s\n", (char *)id->data,