diff --git a/drivers/gpu/nvgpu/common/fifo/tsg.c b/drivers/gpu/nvgpu/common/fifo/tsg.c
index a6f64efe1..9b15ba78f 100644
--- a/drivers/gpu/nvgpu/common/fifo/tsg.c
+++ b/drivers/gpu/nvgpu/common/fifo/tsg.c
@@ -968,6 +968,7 @@ static struct nvgpu_tsg *tsg_gk20a_from_ref(struct nvgpu_ref *ref)
 void nvgpu_tsg_release(struct nvgpu_ref *ref)
 {
 	struct nvgpu_tsg *tsg = tsg_gk20a_from_ref(ref);
+	struct nvgpu_gr_ctx *gr_ctx = tsg->gr_ctx;
 	struct gk20a *g = tsg->g;
 	int err;
 
@@ -977,8 +978,8 @@ void nvgpu_tsg_release(struct nvgpu_ref *ref)
 		return;
 	}
 
-	if ((tsg->gr_ctx != NULL) && (tsg->vm != NULL)) {
-		g->ops.gr.setup.free_gr_ctx(g, tsg->gr_ctx);
+	if ((gr_ctx != NULL) && nvgpu_gr_ctx_get_ctx_initialized(gr_ctx)) {
+		g->ops.gr.setup.free_gr_ctx(g, gr_ctx);
 	}
 
 #ifdef CONFIG_NVGPU_CHANNEL_TSG_CONTROL
diff --git a/drivers/gpu/nvgpu/common/gr/ctx.c b/drivers/gpu/nvgpu/common/gr/ctx.c
index a25793f42..541e5a234 100644
--- a/drivers/gpu/nvgpu/common/gr/ctx.c
+++ b/drivers/gpu/nvgpu/common/gr/ctx.c
@@ -563,6 +563,16 @@ u32 nvgpu_gr_ctx_get_tsgid(struct nvgpu_gr_ctx *gr_ctx)
 	return gr_ctx->tsgid;
 }
 
+void nvgpu_gr_ctx_mark_ctx_initialized(struct nvgpu_gr_ctx *gr_ctx)
+{
+	gr_ctx->ctx_initialized = true;
+}
+
+bool nvgpu_gr_ctx_get_ctx_initialized(struct nvgpu_gr_ctx *gr_ctx)
+{
+	return gr_ctx->ctx_initialized;
+}
+
 #ifdef CONFIG_NVGPU_GRAPHICS
 void nvgpu_gr_ctx_init_graphics_preemption_mode(struct nvgpu_gr_ctx *gr_ctx,
 	u32 graphics_preempt_mode)
diff --git a/drivers/gpu/nvgpu/common/gr/ctx_priv.h b/drivers/gpu/nvgpu/common/gr/ctx_priv.h
index a47228bfd..887fe4731 100644
--- a/drivers/gpu/nvgpu/common/gr/ctx_priv.h
+++ b/drivers/gpu/nvgpu/common/gr/ctx_priv.h
@@ -149,6 +149,8 @@ struct nvgpu_gr_ctx {
 	 */
 	u32 tsgid;
 
+	bool ctx_initialized;
+
 #ifdef CONFIG_NVGPU_SM_DIVERSITY
 	/** SM diversity configuration offset.
 	 * It is valid only if NVGPU_SUPPORT_SM_DIVERSITY support is true.
diff --git a/drivers/gpu/nvgpu/common/gr/gr_setup.c b/drivers/gpu/nvgpu/common/gr/gr_setup.c
index 8dd089626..3b0f0daac 100644
--- a/drivers/gpu/nvgpu/common/gr/gr_setup.c
+++ b/drivers/gpu/nvgpu/common/gr/gr_setup.c
@@ -267,6 +267,8 @@ int nvgpu_gr_setup_alloc_obj_ctx(struct nvgpu_channel *c, u32 class_num,
 	}
 #endif
 
+	nvgpu_gr_ctx_mark_ctx_initialized(gr_ctx);
+
 	nvgpu_mutex_release(&tsg->ctx_init_lock);
 
 	nvgpu_log(g, gpu_dbg_fn | gpu_dbg_gr, "done");
diff --git a/drivers/gpu/nvgpu/common/vgpu/gr/gr_vgpu.c b/drivers/gpu/nvgpu/common/vgpu/gr/gr_vgpu.c
index e968413b6..c50e95720 100644
--- a/drivers/gpu/nvgpu/common/vgpu/gr/gr_vgpu.c
+++ b/drivers/gpu/nvgpu/common/vgpu/gr/gr_vgpu.c
@@ -233,7 +233,9 @@ int vgpu_gr_alloc_obj_ctx(struct nvgpu_channel  *c, u32 class_num, u32 flags)
 
 	err = vgpu_comm_sendrecv(&msg, sizeof(msg), sizeof(msg));
 	err = err ? err : msg.ret;
-	if (err) {
+	if (err == 0) {
+		nvgpu_gr_ctx_mark_ctx_initialized(gr_ctx);
+	} else {
 		nvgpu_err(g, "alloc obj ctx failed err %d", err);
 	}
 	return err;
diff --git a/drivers/gpu/nvgpu/include/nvgpu/gr/ctx.h b/drivers/gpu/nvgpu/include/nvgpu/gr/ctx.h
index 7a703db57..990c68cf6 100644
--- a/drivers/gpu/nvgpu/include/nvgpu/gr/ctx.h
+++ b/drivers/gpu/nvgpu/include/nvgpu/gr/ctx.h
@@ -530,6 +530,28 @@ void nvgpu_gr_ctx_set_tsgid(struct nvgpu_gr_ctx *gr_ctx, u32 tsgid);
  */
 u32 nvgpu_gr_ctx_get_tsgid(struct nvgpu_gr_ctx *gr_ctx);
 
+/**
+ * @brief Mark graphics context initialized.
+ *
+ * @param gr_ctx [in]          Pointer to graphics context struct.
+ *
+ * This function will mark obj_ctx initialized for the gr_ctx by
+ * setting ctx_initialized to true.
+ */
+void nvgpu_gr_ctx_mark_ctx_initialized(struct nvgpu_gr_ctx *gr_ctx);
+
+/**
+ * @brief Get ctx_initialized stored in graphics context structure.
+ *
+ * @param gr_ctx [in]          Pointer to graphics context struct.
+ *
+ * This function will return ctx_initialized stored in #nvgpu_gr_ctx
+ * structure.
+ *
+ * @return ctx initialization status.
+ */
+bool nvgpu_gr_ctx_get_ctx_initialized(struct nvgpu_gr_ctx *gr_ctx);
+
 #ifdef CONFIG_NVGPU_CILP
 bool nvgpu_gr_ctx_get_cilp_preempt_pending(struct nvgpu_gr_ctx *gr_ctx);