fix:missing `output_router_logits` in SwitchTransformers (#30573)

* fix:missing `output_router_logits` in SwitchTransformers * fix whitespace in blank line

fix:missing `output_router_logits` in SwitchTransformers (#30573)
* fix:missing `output_router_logits` in SwitchTransformers * fix whitespace in blank line
a65da83d · Zhan Lu · GitHub · 4ad5adaf · a65da83d
Unverified Commit a65da83d authored May 02, 2024 by Zhan Lu Committed by GitHub May 02, 2024
Hide whitespace changes
Inline Side-by-side

Showing with 3 additions and 0 deletions

src/transformers/models/switch_transformers/modeling_switch_transformers.py ...odels/switch_transformers/modeling_switch_transformers.py +3 -0

No files found.
--- a/src/transformers/models/switch_transformers/modeling_switch_transformers.py
+++ b/src/transformers/models/switch_transformers/modeling_switch_transformers.py
@@ -1721,6 +1721,8 @@ class SwitchTransformersForConditionalGeneration(SwitchTransformersPreTrainedMod

            input_ids = input_ids[:, remove_prefix_length:]

+        output_router_logits = kwargs.get("output_router_logits", True)
+
        return {
            "decoder_input_ids": input_ids,
            "past_key_values": past_key_values,
@@ -1730,6 +1732,7 @@ class SwitchTransformersForConditionalGeneration(SwitchTransformersPreTrainedMod
            "decoder_head_mask": decoder_head_mask,
            "cross_attn_head_mask": cross_attn_head_mask,
            "use_cache": use_cache,
+            "output_router_logits": output_router_logits,
        }

    def prepare_decoder_input_ids_from_labels(self, labels: torch.Tensor):