Update modeling_gemmoe.py

modeling_gemmoe.py

@@ -1086,8 +1086,6 @@ class GemmoeModel(GemmoePreTrainedModel):
         for decoder_layer in self.layers:
             if output_hidden_states:
                 all_hidden_states += (hidden_states,)
-
-            if self.gradient_checkpointing and self.training:
                 layer_outputs = self._gradient_checkpointing_func(
                     decoder_layer.__call__,
                     hidden_states,
@@ -1095,8 +1093,10 @@ class GemmoeModel(GemmoePreTrainedModel):
                     position_ids,
                     past_key_values,
                     output_attentions,
-                    bool(use_cache),
+                    output_router_logits,
+                    use_cache.item() if isinstance(use_cache, torch.Tensor) else use_cache,
                     cache_position,
+                    output_router_logits,
                 )
             else:
                 layer_outputs = decoder_layer(
@@ -1105,7 +1105,8 @@ class GemmoeModel(GemmoePreTrainedModel):
                     position_ids=position_ids,
                     past_key_value=past_key_values,
                     output_attentions=output_attentions,
-                    use_cache=bool(use_cache),
+                    output_router_logits=output_router_logits,
+                    use_cache=use_cache.item() if isinstance(use_cache, torch.Tensor) else use_cache,
                     cache_position=cache_position,
                 )