[Enhance] Support CPU-only train and test in slurm cluster (open-mmla…

…b#189)
GhaSiKey · Jan 4, 2023 · 706e5ee · 706e5ee
1 parent ff137cc
commit 706e5ee
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 7 deletions.
diff --git a/mim/commands/test.py b/mim/commands/test.py
@@ -190,10 +190,8 @@ def test(
 
     # If launcher == "slurm", must have following args
     if launcher == 'slurm':
-        msg = ('If launcher is slurm, '
-               'gpus, gpus-per-node and partition should not be None')
-        flag = (gpus_per_node
-                is not None) and (partition is not None) and (gpus is not None)
+        msg = 'If launcher is slurm, partition should not be None'
+        flag = partition is not None
         assert flag, msg
 
     if not is_installed(package):

diff --git a/mim/commands/train.py b/mim/commands/train.py
@@ -162,9 +162,8 @@ def train(
 
     # If launcher == "slurm", must have following args
     if launcher == 'slurm':
-        msg = ('If launcher is slurm, '
-               'gpus-per-node and partition should not be None')
-        flag = (gpus_per_node is not None) and (partition is not None)
+        msg = 'If launcher is slurm, partition should not be None'
+        flag = partition is not None
         assert flag, msg
 
     if port is None: