diff --git a/drivers/gpu/nvgpu/gk20a/mm_gk20a.c b/drivers/gpu/nvgpu/gk20a/mm_gk20a.c
index 173776ffc..9ab3ae884 100644
--- a/drivers/gpu/nvgpu/gk20a/mm_gk20a.c
+++ b/drivers/gpu/nvgpu/gk20a/mm_gk20a.c
@@ -2379,7 +2379,7 @@ int gk20a_vm_alloc_space(struct gk20a_as_share *as_share,
 
 	va_node->vaddr_start = vaddr_start;
 	va_node->size = (u64)args->page_size * (u64)args->pages;
-	va_node->pgsz_idx = args->page_size;
+	va_node->pgsz_idx = pgsz_idx;
 	INIT_LIST_HEAD(&va_node->va_buffers_list);
 	INIT_LIST_HEAD(&va_node->reserved_va_list);
 
diff --git a/drivers/gpu/nvgpu/gm20b/mm_gm20b.c b/drivers/gpu/nvgpu/gm20b/mm_gm20b.c
index 0f9749f7e..afb04cae4 100644
--- a/drivers/gpu/nvgpu/gm20b/mm_gm20b.c
+++ b/drivers/gpu/nvgpu/gm20b/mm_gm20b.c
@@ -272,19 +272,11 @@ fail:
 }
 
 void gm20b_vm_clear_sparse(struct vm_gk20a *vm, u64 vaddr,
-			       u64 size, u32 pgsz) {
-	int pgsz_idx;
+			       u64 size, u32 pgsz_idx) {
 	u64 vaddr_hi;
 	u32 pde_lo, pde_hi, pde_i;
 
 	gk20a_dbg_fn("");
-	/* determine pagesz idx */
-	for (pgsz_idx = gmmu_page_size_small;
-	     pgsz_idx < gmmu_nr_page_sizes;
-	     pgsz_idx++) {
-		if (gmmu_page_sizes[pgsz_idx] == pgsz)
-			break;
-	}
 	vaddr_hi = vaddr + size - 1;
 	pde_range_from_vaddr_range(vm,
 				   vaddr,
@@ -292,8 +284,8 @@ void gm20b_vm_clear_sparse(struct vm_gk20a *vm, u64 vaddr,
 				   &pde_lo, &pde_hi);
 
 	gk20a_dbg_info("vaddr: 0x%llx, vaddr_hi: 0x%llx, pde_lo: 0x%x, "
-			"pde_hi: 0x%x, pgsz: %d, pde_stride_shift: %d",
-			vaddr, vaddr_hi, pde_lo, pde_hi, pgsz,
+			"pde_hi: 0x%x, pgsz_idx: %d, pde_stride_shift: %d",
+			vaddr, vaddr_hi, pde_lo, pde_hi, pgsz_idx,
 			vm->mm->pde_stride_shift);
 
 	for (pde_i = pde_lo; pde_i <= pde_hi; pde_i++) {