bug fix for megatron patch

a8c5a3cb · Rick Ho · ae10e942 · a8c5a3cb
Commit a8c5a3cb authored Oct 25, 2021 by Rick Ho
Hide whitespace changes
Inline Side-by-side

Showing with 3 additions and 0 deletions

fmoe/megatron/patch.py fmoe/megatron/patch.py +3 -0

No files found.
--- a/fmoe/megatron/patch.py
+++ b/fmoe/megatron/patch.py
 r"""
 Patching some of Megatron-LM's functions to create an MoE model
 """
+import torch
+
+
 def patch_forward_step(forward_step_func):
    r"""
    Patch model's forward_step_func to support balance loss