From 89917dfc03214ac3422a8bd3604dcdc977a285b8 Mon Sep 17 00:00:00 2001
From: Yanyi Liu <wolfsonliu@163.com>
Date: Thu, 5 Aug 2021 09:02:59 +0800
Subject: [PATCH] fix import bug on torch 1.9.0

---
 actnn/actnn/_utils/collate.py    | 5 +++--
 actnn/actnn/_utils/pin_memory.py | 4 +++-
 actnn/actnn/_utils/worker.py     | 2 +-
 actnn/actnn/dataloader.py        | 3 ++-
 4 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/actnn/actnn/_utils/collate.py b/actnn/actnn/_utils/collate.py
index afe5ade..12a54c6 100644
--- a/actnn/actnn/_utils/collate.py
+++ b/actnn/actnn/_utils/collate.py
@@ -7,7 +7,8 @@
 
 import torch
 import re
-from torch._six import container_abcs, string_classes, int_classes
+from torch._six import string_classes
+import collections.abc as container_abcs
 
 np_str_obj_array_pattern = re.compile(r'[SaUO]')
 
@@ -66,7 +67,7 @@ def default_collate(batch):
             return torch.as_tensor(batch)
     elif isinstance(elem, float):
         return torch.tensor(batch, dtype=torch.float64)
-    elif isinstance(elem, int_classes):
+    elif isinstance(elem, int):
         return torch.tensor(batch)
     elif isinstance(elem, string_classes):
         return batch
diff --git a/actnn/actnn/_utils/pin_memory.py b/actnn/actnn/_utils/pin_memory.py
index 055c3cb..5c7928c 100644
--- a/actnn/actnn/_utils/pin_memory.py
+++ b/actnn/actnn/_utils/pin_memory.py
@@ -6,7 +6,9 @@
 """
 
 import torch
-from torch._six import queue, container_abcs, string_classes
+from torch._six import string_classes
+import queue
+import collections.abc as container_abcs
 from . import MP_STATUS_CHECK_INTERVAL
 from torch._utils import ExceptionWrapper
 
diff --git a/actnn/actnn/_utils/worker.py b/actnn/actnn/_utils/worker.py
index fc5f081..762bdf8 100644
--- a/actnn/actnn/_utils/worker.py
+++ b/actnn/actnn/_utils/worker.py
@@ -8,7 +8,7 @@
 import random
 import os
 from collections import namedtuple
-from torch._six import queue
+import queue
 from torch._utils import ExceptionWrapper
 from . import signal_handling, MP_STATUS_CHECK_INTERVAL, IS_WINDOWS
 
diff --git a/actnn/actnn/dataloader.py b/actnn/actnn/dataloader.py
index a212db0..053e260 100644
--- a/actnn/actnn/dataloader.py
+++ b/actnn/actnn/dataloader.py
@@ -13,7 +13,8 @@
 import torch
 import torch.multiprocessing as multiprocessing
 from torch._utils import ExceptionWrapper
-from torch._six import queue, string_classes
+import queue
+from torch._six import string_classes
 
 from torch.utils.data import IterableDataset, Sampler, SequentialSampler, RandomSampler, BatchSampler
 from . import _utils