Remove assertion for ranks for non-dist mode, local dist test

archai/common/apex_utils.py

@@ -66,10 +66,6 @@ class ApexUtils:
                 assert dist.is_initialized()
             assert dist.get_world_size() == self.world_size
             assert dist.get_rank() == self.global_rank
-        else:
-            assert self.world_size == 1
-            assert self.local_rank == 0
-            assert self.global_rank == 0
 
         assert self.world_size >= 1
         assert not self._min_world_size or self.world_size >= self._min_world_size

scripts/dist_main.sh

@@ -15,4 +15,4 @@ echo "*****************************************"
 set -e -o xtrace
 
 
-python -m torch.distributed.launch --nproc_per_node=$gpu_count scripts/main.py --nas.eval.trainer.apex.enabled True $*
+python -m torch.distributed.launch --nproc_per_node=$gpu_count scripts/main.py $*

scripts/misc/dist_local_test.sh

@@ -0,0 +1 @@
+bash scripts/dist_main.sh --no-search  --nas.eval.trainer.apex.enabled True  --nas.eval.final_desc_filename models/darts/darts_genotype.yaml --algos darts