fix for torch2.0.0 suddenly being imposed on us

2023-03-15 19:24:35 +00:00 · 2023-03-15 19:24:35 +00:00 · bea6174a19
commit bea6174a19
parent b253da6e35
1 changed files with 7 additions and 1 deletions
--- a/codes/utils/util.py
+++ b/codes/utils/util.py
@ -22,7 +22,13 @@ from shutil import get_terminal_size
 import scp
 import paramiko
 from torch.utils.checkpoint import checkpoint
-from torch._six import inf
+
+try:
+    # 1.13.1
+    from torch._six import inf
+except Exception as e:
+    # 2.0
+    from torch import inf

 import yaml