From 706e5ee3195789be8c0167e0357f31770b321d2a Mon Sep 17 00:00:00 2001 From: gaoshiqi Date: Wed, 4 Jan 2023 17:34:53 +0800 Subject: [PATCH] [Enhance] Support CPU-only train and test in slurm cluster (#189) --- mim/commands/test.py | 6 ++---- mim/commands/train.py | 5 ++--- 2 files changed, 4 insertions(+), 7 deletions(-) diff --git a/mim/commands/test.py b/mim/commands/test.py index 6d79110..f9fa715 100644 --- 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 index b0273bc..1ce3609 100644 --- 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: