Update megatron/training.py

Commit

4 years ago

Update megatron/training.py Co-authored-by: Stas Bekman <stas00@users.noreply.github.com>