Fixed shape error by unqueezing logits

Aleph-Alpha · Mar 20, 2023 · 4d8d4d7 · 4d8d4d7
1 parent de69f6d
commit 4d8d4d7
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/magma/image_prefix.py b/magma/image_prefix.py
@@ -109,8 +109,8 @@ def forward(
         if self.use_layernorm:
             logits = self.ln(logits)
 
-        # Added for shape mismatch. No longer needed?
-        #if logits.ndim == 2:
-        #    logits = logits.unsqueeze(1)
+        # Added for shape mismatch.
+        if logits.ndim == 2:
+            logits = logits.unsqueeze(1)
 
         return logits