Merge remote-tracking branch 'origin/dev-3.x-folding' into llama_2_demo_transformer_bridge_migration

Author: bryce13950

transformer_lens/model_bridge/bridge.py

@@ -918,8 +918,11 @@ def _ported_forward_pass(
             logits = self.unembed(residual)
             loss = self._calculate_loss(logits, tokens, loss_per_token)
             return logits, loss
+        elif return_type is None:
+            # Return None when explicitly requested
+            return None
         else:
-            # Return final residual
+            # Return final residual for any other return_type
             return residual
 
     def _calculate_loss(self, logits, tokens, loss_per_token=False):
@@ -1678,6 +1681,9 @@ def _handle_return_type(self, logits, input_ids, return_type, loss_per_token):
             shift_logits = logits[:, :-1, :].contiguous()
             loss = F.cross_entropy(shift_logits.view(-1, shift_logits.size(-1)), labels.view(-1))
             return loss, logits
+        elif return_type is None:
+            # Return None when explicitly requested
+            return None
         else:
             return logits
 
@@ -2237,6 +2243,9 @@ def _true_hf_format_forward_pass(
                     shift_logits.view(-1, shift_logits.size(-1)), targets.view(-1), reduction="mean"
                 )
             return (loss, logits)
+        elif return_type is None:
+            # Return None when explicitly requested
+            return None
         else:
             return logits
 
@@ -2421,6 +2430,9 @@ def _hf_format_forward_pass(
                     shift_logits.view(-1, shift_logits.size(-1)), targets.view(-1), reduction="mean"
                 )
             return (loss, logits)
+        elif return_type is None:
+            # Return None when explicitly requested
+            return None
         else:
             return logits
 
@@ -3719,7 +3731,8 @@ def forward(
                 loss = self.loss_fn(logits, input_ids, per_token=loss_per_token)
             return logits, loss
         elif return_type is None:
-            return output
+            # Return None when explicitly requested (don't return output/logits)
+            return None
         else:
             raise ValueError(f"Invalid return_type: {return_type}")