conditionally import amp_C (#1211)

NVIDIA · Nov 10, 2021 · 25bfcb9 · 25bfcb9
1 parent 2205cff
commit 25bfcb9
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/apex/transformer/pipeline_parallel/utils.py b/apex/transformer/pipeline_parallel/utils.py
@@ -20,10 +20,11 @@
 from torch.nn.parallel import DistributedDataParallel
 
 from apex.multi_tensor_apply import multi_tensor_applier
-import amp_C
 from apex.transformer import parallel_state
 from apex.transformer.microbatches import build_num_microbatches_calculator
 from apex.transformer.pipeline_parallel._timers import _Timers
+if multi_tensor_applier.available:
+    import amp_C
 
 
 _GLOBAL_ARGS = None