Fix import for depreciated torch._six

See pytorch/pytorch#94709 (comment)
BINDS-LAB-UMASS · Feb 29, 2024 · 63fe258 · 63fe258
1 parent 3ca676b
commit 63fe258
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 2 deletions.
diff --git a/bindsnet/datasets/collate.py b/bindsnet/datasets/collate.py
@@ -8,7 +8,14 @@
 """
 
 import torch
-from torch._six import container_abcs, string_classes, int_classes
+
+# https://github.com/pytorch/pytorch/pull/94709#issuecomment-1461471006
+try:
+    from torch._six import container_abcs, string_classes, int_classes
+except ImportError:
+    int_classes = int
+    string_classes = str
+    import collections.abc as container_abcs
 
 from torch.utils.data._utils import collate as pytorch_collate
 

diff --git a/bindsnet/pipeline/base_pipeline.py b/bindsnet/pipeline/base_pipeline.py
@@ -2,7 +2,13 @@
 from typing import Tuple, Dict, Any
 
 import torch
-from torch._six import container_abcs, string_classes
+
+# https://github.com/pytorch/pytorch/pull/94709#issuecomment-1461471006
+try:
+    from torch._six import container_abcs, string_classes
+except ImportError:
+    string_classes = str
+    import collections.abc as container_abcs
 
 from ..network import Network
 from ..network.monitors import Monitor