diff --git a/drivers/gpu/nvgpu/common/mm/vm.c b/drivers/gpu/nvgpu/common/mm/vm.c
index 13fd641e2..3d9dd1742 100644
--- a/drivers/gpu/nvgpu/common/mm/vm.c
+++ b/drivers/gpu/nvgpu/common/mm/vm.c
@@ -1094,6 +1094,11 @@ static void __nvgpu_vm_unmap(struct nvgpu_mapped_buf *mapped_buffer,
 	nvgpu_kfree(g, mapped_buffer);
 }
 
+/*
+ * Note: the update_gmmu_lock of the VM that owns this buffer must be locked
+ * before calling nvgpu_ref_put() with this function as the unref function
+ * argument since this can modify the tree of maps.
+ */
 void __nvgpu_vm_unmap_ref(struct nvgpu_ref *ref)
 {
 	struct nvgpu_mapped_buf *mapped_buffer =
diff --git a/drivers/gpu/nvgpu/common/mm/vm_area.c b/drivers/gpu/nvgpu/common/mm/vm_area.c
index 663d89991..5d3b0526a 100644
--- a/drivers/gpu/nvgpu/common/mm/vm_area.c
+++ b/drivers/gpu/nvgpu/common/mm/vm_area.c
@@ -199,7 +199,6 @@ int nvgpu_vm_area_free(struct vm_gk20a *vm, u64 addr)
 		return 0;
 	}
 	nvgpu_list_del(&vm_area->vm_area_list);
-	nvgpu_mutex_release(&vm->update_gmmu_lock);
 
 	nvgpu_log(g, gpu_dbg_map,
 		  "DEL vm_area: pgsz=%#-8x pages=%-9llu "
@@ -219,6 +218,8 @@ int nvgpu_vm_area_free(struct vm_gk20a *vm, u64 addr)
 		nvgpu_ref_put(&buffer->ref, __nvgpu_vm_unmap_ref);
 	}
 
+	nvgpu_mutex_release(&vm->update_gmmu_lock);
+
 	/* if this was a sparse mapping, free the va */
 	if (vm_area->sparse)
 		g->ops.mm.gmmu_unmap(vm,