diff --git a/drivers/gpu/nvgpu/common/gr/gr.c b/drivers/gpu/nvgpu/common/gr/gr.c
index eaddc3cce..696db7faf 100644
--- a/drivers/gpu/nvgpu/common/gr/gr.c
+++ b/drivers/gpu/nvgpu/common/gr/gr.c
@@ -1184,12 +1184,14 @@ void nvgpu_gr_sw_ready(struct gk20a *g, bool enable)
 }
 
 #ifdef CONFIG_NVGPU_NON_FUSA
-/* Wait until GR is initialized */
 void nvgpu_gr_wait_initialized(struct gk20a *g)
 {
 	struct nvgpu_gr *gr = nvgpu_gr_get_cur_instance_ptr(g);
 
-	NVGPU_COND_WAIT(&gr->init_wq, gr->initialized, 0U);
+	/* Wait until GR is initialized */
+	NVGPU_COND_WAIT(&gr->init_wq, gr->initialized ||
+			(nvgpu_is_enabled(g, NVGPU_KERNEL_IS_DYING) ||
+			nvgpu_is_enabled(g, NVGPU_DRIVER_IS_DYING)), 0U);
 }
 #endif
 
diff --git a/drivers/gpu/nvgpu/os/linux/module.c b/drivers/gpu/nvgpu/os/linux/module.c
index c2e32e7b8..aefdfb210 100644
--- a/drivers/gpu/nvgpu/os/linux/module.c
+++ b/drivers/gpu/nvgpu/os/linux/module.c
@@ -61,6 +61,7 @@
 #include <nvgpu/engines.h>
 #include <nvgpu/channel.h>
 #include <nvgpu/gr/gr.h>
+#include <nvgpu/gr/gr_instances.h>
 #include <nvgpu/gr/gr_utils.h>
 #include <nvgpu/pmu/pmu_pstate.h>
 #include <nvgpu/cyclestats_snapshot.h>
@@ -71,6 +72,7 @@
 #include <nvgpu/fb.h>
 #include <nvgpu/nvs.h>
 
+#include "common/gr/gr_priv.h"
 #include "platform_gk20a.h"
 #include "sysfs.h"
 #include "vgpu/vgpu_linux.h"
@@ -147,8 +149,21 @@ static int nvgpu_kernel_shutdown_notification(struct notifier_block *nb,
 	struct nvgpu_os_linux *l = container_of(nb, struct nvgpu_os_linux,
 						nvgpu_reboot_nb);
 	struct gk20a *g = &l->g;
+	struct nvgpu_gr *gr = nvgpu_gr_get_cur_instance_ptr(g);
 
 	nvgpu_set_enabled(g, NVGPU_KERNEL_IS_DYING, true);
+
+	/*
+	 * In rmmod path, when the kernel or GPU driver is
+	 * dying signal the gr wait queue so that the wait
+	 * queue wakes up and further processing happens.
+	 * This is needed to prevent other threads, like
+	 * pmu_pg_task, to go into un-interruptible state.
+	 */
+	if (gr != NULL) {
+		nvgpu_cond_signal(&gr->init_wq);
+	}
+
 	return NOTIFY_DONE;
 }
 
@@ -1590,10 +1605,21 @@ static int gk20a_pm_deinit(struct device *dev)
 
 void nvgpu_start_gpu_idle(struct gk20a *g)
 {
+	struct nvgpu_gr *gr = nvgpu_gr_get_cur_instance_ptr(g);
 	struct nvgpu_os_linux *l = nvgpu_os_linux_from_gk20a(g);
 
 	down_write(&l->busy_lock);
 	nvgpu_set_enabled(g, NVGPU_DRIVER_IS_DYING, true);
+	/*
+	 * In rmmod path, when the kernel or GPU driver is
+	 * dying signal the gr wait queue so that the wait
+	 * queue wakes up and further processing happens.
+	 * This is needed to prevent other threads, like
+	 * pmu_pg_task, to go into un-interruptible state.
+	 */
+	if (gr != NULL) {
+		nvgpu_cond_signal(&gr->init_wq);
+	}
 	/*
 	 * GR SW ready needs to be invalidated at this time with the busy lock
 	 * held to prevent a racing condition on the gr/mm code