fix: use d_att instead of hidden_dim in enable_corr=False fallback for LocalGatedPropagation

model/attention.py

@@ -829,8 +829,9 @@ class LocalGatedPropagation(nn.Module):
             qk = self.correlation_sampler(q, k).view(
                 n, self.num_head, self.window_size * self.window_size, h * w)
         else:
+            # Use self.d_att (key dim) not hidden_dim (value dim) for unfolded keys
             unfolded_k = self.pad_and_unfold(k).view(
-                n * self.num_head, hidden_dim,
+                n * self.num_head, self.d_att,
                 self.window_size * self.window_size, h, w)
             qk = (q.unsqueeze(2) * unfolded_k).sum(dim=1).view(
                 n, self.num_head, self.window_size * self.window_size, h * w)