Update region.py

fix: return shape for decode.

region.py

@@ -71,19 +71,11 @@ def encode_size(size: torch.Tensor, w: nn.Module) -> torch.Tensor:
     return w.size_encoder(fourier_features(size, w.size_features))
 
 
-# region.py
 def decode_size(hidden_state: torch.Tensor, w: nn.Module) -> torch.Tensor:
-    """
-    Returns logits for width & height bins without collapsing batch/seq dims.
+    # Original API expected by moondream.py: shape (2, C) when called on the last hidden state
+    x = mlp(hidden_state, w.size_decoder)  # (..., 2*C)
+    return x.view(2, -1)
 
-    Input  (hidden_state): (..., C)
-    Output: (..., 2, bins)  # keeps all leading dims intact
-    """
-    x = mlp(hidden_state, w.size_decoder)       # (..., size_out_dim)
-    last = x.shape[-1]
-    if last % 2 != 0:
-        raise RuntimeError(f"size_out_dim must be even, got {last}")
-    return x.view(*x.shape[:-1], 2, last // 2)  # (..., 2, bins)