Basic test for segmentation fine-tuning tasks

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

Author: Joao-L-S-Almeida

tests/test_finetune.py

@@ -62,3 +62,16 @@ def test_finetune_bands_str(model_name):
     command_list = ["fit", "-c", f"tests/manufactured-finetune_{model_name}_metrics_from_file.yaml"]
     _ = build_lightning_cli(command_list)
 
+@pytest.mark.parametrize("model_name", ["prithvi_swin_B"])
+def test_finetune_segmentation(model_name):
+
+    model_instance = timm.create_model(model_name)
+
+    state_dict = model_instance.state_dict()
+
+    torch.save(state_dict, os.path.join("tests/", model_name + ".pt"))
+
+    # Running the terratorch CLI
+    command_list = ["fit", "-c", f"tests/manufactured-finetune_{model_name}_segmentation.yaml"]
+    _ = build_lightning_cli(command_list)
+