Fix numpy upgrade to 2.0.0 BUFSIZE import error

microsoft · Jun 18, 2024 · 396c634 · 396c634
1 parent b33873d
commit 396c634
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/deepspeed/autotuning/scheduler.py b/deepspeed/autotuning/scheduler.py
@@ -5,7 +5,11 @@
 
 import copy
 
-from numpy import BUFSIZE
+import numpy
+if numpy.__version__ < '2.0.0':
+    from numpy import BUFSIZE
+else:
+    from numpy._core.umath import BUFSIZE
 import json
 import subprocess
 import sys