diff --git a/drivers/gpu/nvgpu/common/linux/module.c b/drivers/gpu/nvgpu/common/linux/module.c
index cbfe6ad71..bfbe7a58f 100644
--- a/drivers/gpu/nvgpu/common/linux/module.c
+++ b/drivers/gpu/nvgpu/common/linux/module.c
@@ -659,7 +659,7 @@ static void gk20a_pm_shutdown(struct platform_device *pdev)
 	/* Prevent more requests by disabling Runtime PM */
 	__pm_runtime_disable(&pdev->dev, false);
 
-	err = gk20a_wait_for_idle(&pdev->dev);
+	err = gk20a_wait_for_idle(g);
 	if (err) {
 		nvgpu_err(g, "failed to idle GPU, err=%d", err);
 		goto finish;
@@ -844,7 +844,7 @@ void gk20a_driver_start_unload(struct gk20a *g)
 	if (g->is_virtual)
 		return;
 
-	gk20a_wait_for_idle(dev_from_gk20a(g));
+	gk20a_wait_for_idle(g);
 
 	nvgpu_wait_for_deferred_interrupts(g);
 	gk20a_channel_cancel_pending_sema_waits(g);
diff --git a/drivers/gpu/nvgpu/gk20a/gk20a.c b/drivers/gpu/nvgpu/gk20a/gk20a.c
index 7bd6acfdf..c50d800fb 100644
--- a/drivers/gpu/nvgpu/gk20a/gk20a.c
+++ b/drivers/gpu/nvgpu/gk20a/gk20a.c
@@ -362,9 +362,8 @@ int gk20a_can_busy(struct gk20a *g)
 	return 1;
 }
 
-int gk20a_wait_for_idle(struct device *dev)
+int gk20a_wait_for_idle(struct gk20a *g)
 {
-	struct gk20a *g = get_gk20a(dev);
 	int wait_length = 150; /* 3 second overall max wait. */
 	int target_usage_count = 0;
 
diff --git a/drivers/gpu/nvgpu/gk20a/gk20a.h b/drivers/gpu/nvgpu/gk20a/gk20a.h
index aeba1fcd6..4517f6e03 100644
--- a/drivers/gpu/nvgpu/gk20a/gk20a.h
+++ b/drivers/gpu/nvgpu/gk20a/gk20a.h
@@ -1455,7 +1455,7 @@ int __gk20a_do_idle(struct gk20a *g, bool force_reset);
 int __gk20a_do_unidle(struct gk20a *g);
 
 int gk20a_can_busy(struct gk20a *g);
-int gk20a_wait_for_idle(struct device *dev);
+int gk20a_wait_for_idle(struct gk20a *g);
 
 #define NVGPU_GPU_ARCHITECTURE_SHIFT 4