Update model.py

model.py

@@ -45,7 +45,7 @@ class minGRU(nn.Module):
         gate   = gate.sigmoid()
         out    = torch.lerp(h_prev, hidden, gate) if h_prev is not None else (hidden * gate)
 
-        h_next = out[:, -1]
+        h_next = out[:, -1:]
         out    = self.out_proj(out)
 
         return out, h_next
@@ -70,7 +70,7 @@ class minGRU(nn.Module):
         out = heinsen_associative_scan_log(log_coeffs, log_values)
         out = out[:, -seq_len:]
         
-        h_next = out[:, -1]
+        h_next = out[:, -1:]
         out    = self.out_proj(out)
 
         return out, h_next