diff --git a/drivers/gpu/nvgpu/common/device.c b/drivers/gpu/nvgpu/common/device.c
index 47da59b89..13a346f06 100644
--- a/drivers/gpu/nvgpu/common/device.c
+++ b/drivers/gpu/nvgpu/common/device.c
@@ -116,6 +116,13 @@ void nvgpu_device_cleanup(struct gk20a *g)
 	u32 i;
 	struct nvgpu_list_node *devlist;
 
+	/*
+	 * Make unit testing a bit easier.
+	 */
+	if (g->devs == NULL) {
+		return;
+	}
+
 	for (i = 0; i < NVGPU_MAX_DEVTYPE; i++) {
 		devlist = &g->devs->devlist_heads[i];
 
diff --git a/drivers/gpu/nvgpu/common/init/nvgpu_init.c b/drivers/gpu/nvgpu/common/init/nvgpu_init.c
index 13a15b3b7..fd37d95ba 100644
--- a/drivers/gpu/nvgpu/common/init/nvgpu_init.c
+++ b/drivers/gpu/nvgpu/common/init/nvgpu_init.c
@@ -326,8 +326,6 @@ int nvgpu_prepare_poweroff(struct gk20a *g)
 	}
 #endif
 
-	nvgpu_device_cleanup(g);
-
 	/* Disable GPCPLL */
 	if (g->ops.clk.suspend_clk_support != NULL) {
 		g->ops.clk.suspend_clk_support(g);
@@ -868,6 +866,8 @@ static void gk20a_free_cb(struct nvgpu_ref *refcount)
 		g->ops.ltc.ltc_remove_support(g);
 	}
 
+	nvgpu_device_cleanup(g);
+
 	nvgpu_sw_quiesce_remove_support(g);
 
 	if (g->gfree != NULL) {