Merge pull request #2 from Quentin-Anthony/shape-fix

Fixed shape error by unqueezing logits
Aleph-Alpha · Mar 23, 2023 · d793893 · d793893
2 parents 27f81d6 + 4d8d4d7
commit d793893
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