diff --git a/drivers/staging/imx-drm/imx-drm-core.c b/drivers/staging/imx-drm/imx-drm-core.c
index 6bd015ac9d683474a034924f8ffec3e752e1d382..0507b662ae40d331ca257b751c2fb95af9cafe41 100644
--- a/drivers/staging/imx-drm/imx-drm-core.c
+++ b/drivers/staging/imx-drm/imx-drm-core.c
@@ -444,6 +444,8 @@ static int imx_drm_driver_load(struct drm_device *drm, unsigned long flags)
 	if (!imx_drm_device_get())
 		ret = -EINVAL;
 
+	platform_set_drvdata(drm->platformdev, drm);
+
 	ret = 0;
 
 err_init:
@@ -829,7 +831,7 @@ static int imx_drm_platform_probe(struct platform_device *pdev)
 
 static int imx_drm_platform_remove(struct platform_device *pdev)
 {
-	drm_platform_exit(&imx_drm_driver, pdev);
+	drm_put_dev(platform_get_drvdata(pdev));
 
 	return 0;
 }