diff --git a/bonsai/models/dinov3/modeling.py b/bonsai/models/dinov3/modeling.py
index 9ab3b0df..22ed2a43 100644
--- a/bonsai/models/dinov3/modeling.py
+++ b/bonsai/models/dinov3/modeling.py
@@ -1,4 +1,5 @@
 import dataclasses
+from typing import Literal
 
 import jax
 import jax.numpy as jnp
@@ -13,7 +14,7 @@ class ModelConfig:
     intermediate_size: int = 1536
     num_hidden_layers: int = 12
     num_attention_heads: int = 6
-    hidden_act: str = "gelu"
+    hidden_act: Literal["gelu", "silu"] = "gelu"
     layer_norm_eps: float = 1e-5
     rope_theta: float = 100.0
     image_size: int = 224
@@ -127,12 +128,12 @@ def __init__(self, config: ModelConfig):
         self.base = config.rope_theta
         self.head_dim = config.hidden_size // config.num_attention_heads
         self.num_patches_h = config.image_size // config.patch_size[0]
-        self.num_patches_w = config.image_size // config.patch_size[0]
+        self.num_patches_w = config.image_size // config.patch_size[1]
 
     def __call__(self, pixel_values: Array) -> tuple[Array, Array]:
         _, _, height, width = pixel_values.shape
         num_patches_h = height // self.config.patch_size[0]
-        num_patches_w = width // self.config.patch_size[0]
+        num_patches_w = width // self.config.patch_size[1]
 
         coords_h = jnp.arange(0.5, num_patches_h, dtype=jnp.float32) / num_patches_h  # [H]
         coords_w = jnp.arange(0.5, num_patches_w, dtype=jnp.float32) / num_patches_w  # [W]
@@ -157,7 +158,7 @@ def __init__(self, config: ModelConfig):
         self.lambda1 = nnx.Param(jnp.full((config.hidden_size,), config.layerscale_value, dtype=jnp.float32))
 
     def __call__(self, x: Array) -> Array:
-        return x * self.lambda1
+        return x * self.lambda1[...]
 
 
 def rotate_half(x: Array) -> Array:
diff --git a/bonsai/models/dinov3/tests/test_outputs_dinov3.py b/bonsai/models/dinov3/tests/test_outputs_dinov3.py
index 9c84009a..031a64db 100644
--- a/bonsai/models/dinov3/tests/test_outputs_dinov3.py
+++ b/bonsai/models/dinov3/tests/test_outputs_dinov3.py
@@ -96,5 +96,25 @@ def test_pooled_output_embeddings(self):
         np.testing.assert_allclose(jy, ty.detach().cpu().numpy(), rtol=1e-5, atol=2e-2)
 
 
+class TestRopePositionEmbedding(absltest.TestCase):
+    def test_non_square_patch_size_uses_width_patch_dimension(self):
+        config = model_lib.ModelConfig(
+            patch_size=(16, 8),
+            hidden_size=64,
+            num_attention_heads=1,
+            image_size=64,
+        )
+        rope = model_lib.Dinov3ViTRopePositionEmbedding(config)
+
+        x = jnp.zeros((1, 3, 64, 40), dtype=jnp.float32)
+        cos, sin = rope(x)
+
+        _, _, height, width = x.shape
+        expected_num_patches = (height // config.patch_size[0]) * (width // config.patch_size[1])
+        expected_head_dim = config.hidden_size // config.num_attention_heads
+        self.assertEqual(cos.shape, (expected_num_patches, expected_head_dim))
+        self.assertEqual(sin.shape, (expected_num_patches, expected_head_dim))
+
+
 if __name__ == "__main__":
     absltest.main()