[moe] fix mixtral forward default value (#5329)

956b561b · Hongxin Liu · ver217 · b60be18d · 956b561b
Commit 956b561b authored Jan 30, 2024 by Hongxin Liu Committed by ver217 Feb 07, 2024
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

applications/ColossalMoE/colossal_moe/models/mixtral_policy.py ...cations/ColossalMoE/colossal_moe/models/mixtral_policy.py +1 -1

No files found.
--- a/applications/ColossalMoE/colossal_moe/models/mixtral_policy.py
+++ b/applications/ColossalMoE/colossal_moe/models/mixtral_policy.py
@@ -437,7 +437,7 @@ class MixtralPipelineForwards:
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
-        output_router_logits: Optional[bool] = True,
+        output_router_logits: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        stage_manager: Optional[PipelineStageManager] = None,
        hidden_states: Optional[torch.FloatTensor] = None,