diff --git a/drivers/gpu/nvgpu/gk20a/gk20a.c b/drivers/gpu/nvgpu/gk20a/gk20a.c
index c606e0272..93852c171 100644
--- a/drivers/gpu/nvgpu/gk20a/gk20a.c
+++ b/drivers/gpu/nvgpu/gk20a/gk20a.c
@@ -1166,7 +1166,10 @@ static void gk20a_pm_shutdown(struct platform_device *pdev)
 #endif
 
 	/* Be ready for rail-gate after this point */
-	gk20a_pm_prepare_poweroff(&pdev->dev);
+	if (gk20a_gpu_is_virtual(pdev))
+		vgpu_pm_prepare_poweroff(&pdev->dev);
+	else
+		gk20a_pm_prepare_poweroff(&pdev->dev);
 }
 
 #ifdef CONFIG_PM