Update modeling_gemmoe.py

modeling_gemmoe.py

@@ -1235,16 +1235,13 @@ class GemmoeForCausalLM(GemmoePreTrainedModel):
         logits = self.lm_head(hidden_states)
         logits = logits.float()
 
-            # Handle unused parameters
-    if self.training:
-        for expert in self.model.block_sparse_moe.experts:
-            for param in expert.parameters():
-                if param.requires_grad and param.grad is None:
-                    param.grad = torch.zeros_like(param)
-
-    loss = None
-    if labels is not None:
-        
+        # Handle unused parameters
+        if self.training:
+            for expert in self.model.layers[-1].block_sparse_moe.experts:
+                for param in expert.parameters():
+                    if param.requires_grad and param.grad is None:
+                        param.grad = torch.zeros_like(param)
+
         loss = None
         if labels is not None:
             shift_logits = logits[..., :-1, :].contiguous()