modify batch and fix tensor issue

myolmoe/modeling_myolmoe.py

@@ -1065,7 +1065,7 @@ class MyOlmoeForCausalLM(OlmoePreTrainedModel, GenerationMixin):
                 output = (aux_loss,) + output
             return (loss,) + output if loss is not None else output
         #
-        total_small_expert_loss = 0
+        total_small_expert_loss = torch.tensor(0.0, device=logits.device)
         for layer_output in outputs:
             if len(layer_output) > 1 and isinstance(layer_output[1], torch.Tensor):
                 total_small_expert_loss += layer_output[1]

scripts/train.py

@@ -72,7 +72,7 @@ def main():
     # Training arguments
     training_args = TrainingArguments(
         output_dir="./output",
-        per_device_train_batch_size=2,
+        per_device_train_batch_size=16,
         gradient_accumulation_steps=8,
         learning_rate=1e-5,
         num_train_epochs=1,