修正 PyTorch 1.8 兼容性问题 / PyTorch 1.8 compatability fix (HIT-SCIR#519)

jrjdr · Jun 28, 2021 · 78e760d · 78e760d
1 parent 058c374
commit 78e760d
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 2 deletions.
diff --git a/ltp/data/utils/collate.py b/ltp/data/utils/collate.py
@@ -3,9 +3,16 @@
 # Author: Yunlong Feng <[email protected]>
 
 import torch
-from torch._six import int_classes, string_classes, container_abcs
+from torch._six import string_classes
 from torch.utils.data._utils.collate import np_str_obj_array_pattern, default_collate_err_msg_format
 
+_TORCH_MAJOR, _TORCH_MINOR = map(int, torch.__version__.split('.')[0:2])
+
+if _TORCH_MAJOR < 1 or (_TORCH_MAJOR == 1 and _TORCH_MINOR < 8):
+    from torch._six import int_classes, container_abcs
+else:
+    int_classes = int
+    import collections.abc as container_abcs
 
 def collate(batch):
     r"""Puts each data field into a tensor with outer dimension batch size"""

diff --git a/ltp/utils/convertor.py b/ltp/utils/convertor.py
@@ -3,7 +3,13 @@
 # Author: Yunlong Feng <[email protected]>
 
 import torch
-from torch._six import container_abcs
+
+_TORCH_MAJOR, _TORCH_MINOR = map(int, torch.__version__.split('.')[0:2])
+
+if _TORCH_MAJOR < 1 or (_TORCH_MAJOR == 1 and _TORCH_MINOR < 8):
+    from torch._six import container_abcs
+else:
+    import collections.abc as container_abcs
 
 
 def map2device(batch, device=torch.device('cpu')):