Update code for pytorch 1.9.0

timesformer/datasets/multigrid_helper.py

@@ -3,9 +3,16 @@
 """Helper functions for multigrid training."""
 
 import numpy as np
-from torch._six import int_classes as _int_classes
+import torch
 from torch.utils.data.sampler import Sampler
 
+TORCH_MAJOR = int(torch.__version__.split('.')[0])
+TORCH_MINOR = int(torch.__version__.split('.')[1])
+
+if TORCH_MAJOR == 1 and TORCH_MINOR < 8:
+    from torch._six import int_classes as _int_classes
+else:
+    _int_classes = int
 
 class ShortCycleBatchSampler(Sampler):
     """

timesformer/models/resnet_helper.py

@@ -12,7 +12,6 @@
 from einops import rearrange, reduce, repeat
 import torch.nn.functional as F
 from torch.nn.modules.module import Module
-from torch.nn.modules.linear import _LinearWithBias
 from torch.nn.modules.activation import MultiheadAttention
 
 import numpy as np

timesformer/models/vit_utils.py

@@ -11,7 +11,14 @@
 from timesformer.models.helpers import load_pretrained
 from .build import MODEL_REGISTRY
 from itertools import repeat
-from torch._six import container_abcs
+
+TORCH_MAJOR = int(torch.__version__.split('.')[0])
+TORCH_MINOR = int(torch.__version__.split('.')[1])
+
+if TORCH_MAJOR == 1 and TORCH_MINOR < 8:
+    from torch._six import container_abcs
+else:
+    import collections.abc as container_abcs
 
 DEFAULT_CROP_PCT = 0.875
 IMAGENET_DEFAULT_MEAN = (0.485, 0.456, 0.406)