diff --git a/drivers/gpu/nvgpu/gk20a/gr_gk20a.c b/drivers/gpu/nvgpu/gk20a/gr_gk20a.c
index 6c885b599..0250e97e6 100644
--- a/drivers/gpu/nvgpu/gk20a/gr_gk20a.c
+++ b/drivers/gpu/nvgpu/gk20a/gr_gk20a.c
@@ -1853,11 +1853,11 @@ int gr_gk20a_update_hwpm_ctxsw_mode(struct gk20a *g,
 
 	data |= pm_ctx->pm_mode;
 
-	nvgpu_mem_wr(g, gr_mem, ctxsw_prog_main_image_pm_o(), data);
-
 	if (ctxheader->gpu_va) {
+		nvgpu_mem_wr(g, ctxheader, ctxsw_prog_main_image_pm_o(), data);
 		g->ops.gr.write_pm_ptr(g, ctxheader, virt_addr);
 	} else {
+		nvgpu_mem_wr(g, gr_mem, ctxsw_prog_main_image_pm_o(), data);
 		g->ops.gr.write_pm_ptr(g, gr_mem, virt_addr);
 	}