1 files changed, 5 insertions, 0 deletions

diff --git a/drivers/gpu/nvgpu/common/mm/gmmu.c b/drivers/gpu/nvgpu/common/mm/gmmu.c
index 1eed3a3b..d6aaf8cd 100644
--- a/drivers/gpu/nvgpu/common/mm/gmmu.c
+++ b/drivers/gpu/nvgpu/common/mm/gmmu.c
@@ -828,6 +828,11 @@ static int __nvgpu_locate_pte(struct gk20a *g, struct vm_gk20a *vm,
                 if (!pd_next->mem)
                         return -EINVAL;
 
+                attrs->pgsz = l->get_pgsz(g, pd, pd_idx);
+
+                if (attrs->pgsz >= gmmu_nr_page_sizes)
+                        return -EINVAL;
+
                 return __nvgpu_locate_pte(g, vm, pd_next,
                                           vaddr, lvl + 1, attrs,
                                           data, pd_out, pd_idx_out,