Fix more tests

ekut-es · Jan 11, 2024 · ccfacce · ccfacce
1 parent f5a9922
commit ccfacce
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/test/test_train.py b/test/test_train.py
@@ -104,7 +104,7 @@ def test_datasets(model, dataset, split):
     "model", ["timm_resnet50", "timm_efficientnet_lite1", "timm_focalnet_base_srf"]
 )
 def test_2d(model):
-    command_line = f"hannah-train module=image_classifier dataset=fake2d features=identity trainer.devices=[0] model={model}  trainer.fast_dev_run=true scheduler.max_lr=2.5 module.batch_size=2"
+    command_line = f"hannah-train module=image_classifier dataset=fake2d features=identity trainer.devices=1 model={model}  trainer.fast_dev_run=true scheduler.max_lr=2.5 module.batch_size=2"
     subprocess.run(command_line, shell=True, check=True, cwd=topdir)
 
 
@@ -113,7 +113,7 @@ def test_2d(model):
     "model", ["timm_resnet50", "timm_efficientnet_lite1", "timm_resnet18"]
 )
 def test_cifar_2d(model):
-    command_line = f"hannah-train module=image_classifier dataset=cifar10 features=identity trainer.devices=[0] model={model}  trainer.fast_dev_run=true scheduler.max_lr=2.5 module.batch_size=2"
+    command_line = f"hannah-train module=image_classifier dataset=cifar10 features=identity trainer.devices=1 model={model}  trainer.fast_dev_run=true scheduler.max_lr=2.5 module.batch_size=2"
     subprocess.run(command_line, shell=True, check=True, cwd=topdir)