From bea6174a19ac77ce15921173127bdc17e0e4018b Mon Sep 17 00:00:00 2001
From: mrq <mrq@ecker.tech>
Date: Wed, 15 Mar 2023 19:24:35 +0000
Subject: [PATCH] fix for torch2.0.0 suddenly being imposed on us

---
 codes/utils/util.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/codes/utils/util.py b/codes/utils/util.py
index d2077701..90e4cfdd 100644
--- 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