Checking if the fine-tune command returned 0

Signed-off-by: João Lucas de Sousa Almeida <[email protected]>

Author: Joao-L-S-Almeida

tests/manufactured-finetune_prithvi_swin_B.yaml

@@ -98,7 +98,7 @@ model:
       pretrained: true
       backbone: prithvi_swin_B
       backbone_pretrained_cfg_overlay:
-        file: /tmp/prithvi_swin_B.pt
+        file: tests/prithvi_swin_B.pt
       backbone_drop_path_rate: 0.3
       # backbone_window_size: 8
       decoder_channels: 256

tests/manufactured-finetune_prithvi_swin_L.yaml

@@ -98,7 +98,7 @@ model:
       pretrained: true
       backbone: prithvi_swin_L
       backbone_pretrained_cfg_overlay:
-        file: /tmp/prithvi_swin_L.pt
+        file: tests/prithvi_swin_L.pt
       backbone_drop_path_rate: 0.3
       # backbone_window_size: 8
       decoder_channels: 64

tests/manufactured-finetune_prithvi_vit_100.yaml

@@ -98,7 +98,7 @@ model:
       pretrained: true
       backbone: prithvi_vit_100
       backbone_pretrained_cfg_overlay:
-        file: /tmp/prithvi_vit_100.pt
+        file: tests/prithvi_vit_100.pt
       backbone_drop_path_rate: 0.3
       num_frames: 1
       # backbone_window_size: 8

tests/manufactured-finetune_prithvi_vit_300.yaml

@@ -98,7 +98,7 @@ model:
       pretrained: true
       backbone: prithvi_vit_300
       backbone_pretrained_cfg_overlay:
-        file: /tmp/prithvi_vit_300.pt
+        file: tests/prithvi_vit_300.pt
       backbone_drop_path_rate: 0.3
       # backbone_window_size: 8
       decoder_channels: 64

tests/test_finetune.py

@@ -34,12 +34,12 @@ def test_finetune_multiple_backbones(model_name):
 
     torch.save(filtered_state_dict, os.path.join("tests/", model_name + ".pt"))
 
-    try:
-        # Running the terratorch CLI
-        command_str = f"terratorch fit -c tests/manufactured-finetune_{model_name}.yaml" 
+    # Running the terratorch CLI
+    command_str = f"terratorch fit -c tests/manufactured-finetune_{model_name}.yaml" 
+
+    command_out = subprocess.run(command_str, shell=True) 
+
+    assert not command_out.returncode
 
-        subprocess.run(command_str, shell=True) 
-    except:
-        Exception("Fine-tuning cannot be executed.")