diff --git a/src/maxdiffusion/models/wan/autoencoder_kl_wan.py b/src/maxdiffusion/models/wan/autoencoder_kl_wan.py
index 999f5234b..7d3bee300 100644
--- a/src/maxdiffusion/models/wan/autoencoder_kl_wan.py
+++ b/src/maxdiffusion/models/wan/autoencoder_kl_wan.py
@@ -48,6 +48,21 @@ def _update_cache(cache, idx, value):
   return cache[:idx] + (value,) + cache[idx + 1 :]
 
 
+def _spatial_sharding(mesh: Optional[jax.sharding.Mesh], shape: Sequence[int]) -> Optional[NamedSharding]:
+  """Returns the sharding of a (N, D, H, W, C) array over the vae_spatial axis.
+
+  Shards the width, or the height if only the height is divisible by
+  vae_spatial. If neither is, the width is still sharded and XLA pads it.
+  """
+  if mesh is None or "vae_spatial" not in mesh.shape:
+    return None
+  _, _, height, width, _ = shape
+  num_shards = mesh.shape["vae_spatial"]
+  if width % num_shards != 0 and height % num_shards == 0:
+    return NamedSharding(mesh, P("redundant", None, "vae_spatial", None, None))
+  return NamedSharding(mesh, P("redundant", None, None, "vae_spatial", None))
+
+
 # Helper to ensure kernel_size, stride, padding are tuples of 3 integers
 def _canonicalize_tuple(x: Union[int, Sequence[int]], rank: int, name: str) -> Tuple[int, ...]:
   """Canonicalizes a value to a tuple of integers."""
@@ -121,7 +136,7 @@ def __init__(
 
   def __call__(self, x: jax.Array, cache_x: Optional[jax.Array] = None, idx=-1) -> jax.Array:
     if self.mesh is not None and "vae_spatial" in self.mesh.shape:
-      spatial_sharding = NamedSharding(self.mesh, P("redundant", None, None, "vae_spatial", None))
+      spatial_sharding = _spatial_sharding(self.mesh, x.shape)
       if spatial_sharding is not None:
         x = jax.lax.with_sharding_constraint(x, spatial_sharding)
 
@@ -1169,12 +1184,6 @@ def _encode(self, x: jax.Array, feat_cache: AutoencoderKLWanCache):
     iter_ = 1 + ((t - 1 + CHUNK_SIZE - 1) // CHUNK_SIZE) if t > 1 else 1
     enc_feat_map = feat_cache._enc_feat_map
 
-    spatial_sharding = (
-        NamedSharding(self.mesh, P("redundant", None, None, "vae_spatial", None))
-        if self.mesh is not None and "vae_spatial" in self.mesh.shape
-        else None
-    )
-
     def finalize(out, enc_feat_map):
       feat_cache._enc_feat_map = enc_feat_map
       enc = self.quant_conv(out)
@@ -1184,6 +1193,7 @@ def finalize(out, enc_feat_map):
     with jax.named_scope("AutoencoderKLWan_encode_chunk_0"):
       chunk_0 = x[:, :1, ...]
       out_0, enc_feat_map, _ = self.encoder(chunk_0, feat_cache=enc_feat_map, feat_idx=0)
+      spatial_sharding = _spatial_sharding(self.mesh, out_0.shape)
       if spatial_sharding is not None:
         out_0 = jax.lax.with_sharding_constraint(out_0, spatial_sharding)
 
@@ -1280,11 +1290,7 @@ def _decode(
     x = self.post_quant_conv(z)
 
     dec_feat_map = feat_cache._feat_map
-    spatial_sharding = (
-        NamedSharding(self.mesh, P("redundant", None, None, "vae_spatial", None))
-        if self.mesh is not None and "vae_spatial" in self.mesh.shape
-        else None
-    )
+    spatial_sharding = _spatial_sharding(self.mesh, x.shape)
 
     # First chunk (i=0)
     with jax.named_scope("AutoencoderKLWan_decode_chunk_0"):