diff --git a/drivers/gpu/nvgpu/gk20a/gr_gk20a.c b/drivers/gpu/nvgpu/gk20a/gr_gk20a.c
index 69fe4901a..976e91711 100644
--- a/drivers/gpu/nvgpu/gk20a/gr_gk20a.c
+++ b/drivers/gpu/nvgpu/gk20a/gr_gk20a.c
@@ -2830,10 +2830,13 @@ int gk20a_alloc_obj_ctx(struct channel_gk20a  *c,
 	} else {
 		if (!tsg->tsg_gr_ctx) {
 			tsg->vm = c->vm;
+			gk20a_vm_get(tsg->vm);
 			err = gr_gk20a_alloc_tsg_gr_ctx(g, tsg);
 			if (err) {
 				gk20a_err(dev_from_gk20a(g),
 					"fail to allocate TSG gr ctx buffer");
+				gk20a_vm_put(tsg->vm);
+				tsg->vm = NULL;
 				goto out;
 			}
 		}
diff --git a/drivers/gpu/nvgpu/gk20a/tsg_gk20a.c b/drivers/gpu/nvgpu/gk20a/tsg_gk20a.c
index e7b708941..849c3f6ab 100644
--- a/drivers/gpu/nvgpu/gk20a/tsg_gk20a.c
+++ b/drivers/gpu/nvgpu/gk20a/tsg_gk20a.c
@@ -161,8 +161,10 @@ static void gk20a_tsg_release(struct kref *ref)
 		gr_gk20a_free_tsg_gr_ctx(tsg);
 		tsg->tsg_gr_ctx = NULL;
 	}
-	if (tsg->vm)
+	if (tsg->vm) {
+		gk20a_vm_put(tsg->vm);
 		tsg->vm = NULL;
+	}
 
 	release_used_tsg(&g->fifo, tsg);