diff --git a/model/riscv_vmem.sail b/model/riscv_vmem.sail
index 2d1e5d162..754907382 100644
--- a/model/riscv_vmem.sail
+++ b/model/riscv_vmem.sail
@@ -107,7 +107,7 @@ function pt_walk(sv_params,
   let mem_result = mem_read_priv(Read(Data),              // AccessType
                                  Supervisor,              // Privilege
                                  pte_phys_addr,
-                                 8,                       // atom (8)
+                                 2 ^ sv_params.log_pte_size_bytes,
                                  false,                   // aq
                                  false,                   // rl
                                  false);                  // res
@@ -115,6 +115,9 @@ function pt_walk(sv_params,
   match mem_result {
     MemException(_) => PTW_Failure(PTW_Access(), ext_ptw),
     MemValue(pte)   => {
+      // Extend to 64 bits even on RV32 for simplicity.
+      let pte : bits(64) = zero_extend(pte);
+
       let pte_flags = Mk_PTE_Flags(pte[7 .. 0]);
       if pte_is_invalid(pte_flags) then
         PTW_Failure(PTW_Invalid_PTE(), ext_ptw)