fix import container_abcs issue (#1049)

ShoufaChen · web-flow · commit a78ccf0b3e3f · 2021-02-09T23:12:42.000-07:00
* copy-paste friendly * fix import container_abcs issue Nightly PyTorch has removed `container_abcs` from `torch._six`. pytorch/pytorch@58eb233#diff-b3c160475f0fbe8ad50310f92d3534172ba98203387a962b7dc8f4a23b15cf4dL35 * fix import container_abcs issue Nightly PyTorch has removed `container_abcs` from `torch._six`. pytorch/pytorch@58eb233#diff-b3c160475f0fbe8ad50310f92d3534172ba98203387a962b7dc8f4a23b15cf4dL35 * keep existing for pytorch1.7 and earlier
diff --git a/README.md b/README.md
@@ -122,14 +122,14 @@ See the [Docker example folder](https://github.com/NVIDIA/apex/tree/master/examp
 For performance and full functionality, we recommend installing Apex with
 CUDA and C++ extensions via
 ```
-$ git clone https://github.com/NVIDIA/apex
-$ cd apex
-$ pip install -v --disable-pip-version-check --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./
+git clone https://github.com/NVIDIA/apex
+cd apex
+pip install -v --disable-pip-version-check --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./
 ```
 
 Apex also supports a Python-only build (required with Pytorch 0.4) via
 ```
-$ pip install -v --disable-pip-version-check --no-cache-dir ./
+pip install -v --disable-pip-version-check --no-cache-dir ./
 ```
 A Python-only build omits:
 - Fused kernels required to use `apex.optimizers.FusedAdam`.
diff --git a/apex/amp/_amp_state.py b/apex/amp/_amp_state.py
@@ -8,10 +8,11 @@
 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:
+
+if TORCH_MAJOR == 1 and TORCH_MINOR < 8:
     from torch._six import container_abcs
+else:
+    import collections.abc as container_abcs
 
 
 class AmpState(object):
