use cat instead of creating new tensor

fmoe/megatron/balance.py

@@ -97,12 +97,7 @@ def patch_forward_step(forward_step_func):
         loss_list = [l.mlp.gate.get_loss(clear=False) for l in model.language_model.transformer.layers]
         (loss, state_dict), bal_loss = (
             output,
-            (
-                torch.tensor(
-                    loss_list, device=loss_list[0].device
-                ).mean()
-                * args.balance_loss_weight
-            ).float(),
+            torch.cat(loss_list).mean() * args.balance_loss_weight
         )
 
         # avarage across world group