use cat instead of creating new tensor

913d7127 · Rick Ho · 295a615a · 913d7127
Commit 913d7127 authored Jun 17, 2021 by Rick Ho
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 6 deletions

fmoe/megatron/balance.py fmoe/megatron/balance.py +1 -6

No files found.
--- a/fmoe/megatron/balance.py
+++ b/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.cat(loss_list).mean() * args.balance_loss_weight
-                torch.tensor(
-                    loss_list, device=loss_list[0].device
-                ).mean()
-                * args.balance_loss_weight
-            ).float(),
        )
        # avarage across world group