Fix for #186

apex/amp/_amp_state.py

@@ -5,6 +5,15 @@
 import os
 import torch
 
+TORCH_MAJOR = int(torch.__version__.split('.')[0])
+TORCH_MINOR = int(torch.__version__.split('.')[1])
+
+if TORCH_MAJOR == 0:
+    import collections.abc as container_abcs
+else:
+    from torch._six import container_abcs
+
+
 class AmpState(object):
     def __init__(self):
         self.hard_override=False

apex/amp/_initialize.py

 import torch
-from torch._six import container_abcs, string_classes
+from torch._six import string_classes
 import functools
-from ._amp_state import _amp_state, warn_or_err
+from ._amp_state import _amp_state, warn_or_err, container_abcs
 from .handle import disable_casts
 from .scaler import LossScaler
 from apex.fp16_utils import convert_network