Update build.py (#692)

JDAI-CV · Apr 5, 2023 · 817c748 · 817c748
1 parent 39887a1
commit 817c748
Showing 1 changed file with 8 additions and 1 deletion.
diff --git a/fastreid/data/build.py b/fastreid/data/build.py
@@ -8,7 +8,14 @@
 import os
 
 import torch
-from torch._six import string_classes
+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 string_classes
+else:
+    string_classes = str
+
 from collections import Mapping
 
 from fastreid.config import configurable