addressed jareds comments

b6e73188 · Mohammad · 99410264 · b6e73188
Commit b6e73188 authored Apr 15, 2020 by Mohammad
Hide whitespace changes
Inline Side-by-side

Showing with 4 additions and 0 deletions

megatron/fp16/loss_scaler.py megatron/fp16/loss_scaler.py +4 -0

No files found.
--- a/megatron/fp16/loss_scaler.py
+++ b/megatron/fp16/loss_scaler.py
@@ -14,6 +14,10 @@
 # limitations under the License.

 import torch
+
+from apex.multi_tensor_apply import multi_tensor_applier
+import amp_C
+
 from megatron import mpu

 # item() is a recent addition, so this helps with backward compatibility.