Fix inference code

Signed-off-by: Benedikt Blumenstiel <[email protected]>

config.yaml

@@ -6,18 +6,7 @@ trainer:
   devices: auto
   num_nodes: 1
   precision: 16-mixed
-  logger:
-    class_path: lightning.pytorch.loggers.TensorBoardLogger
-    init_args:
-      save_dir: /dccstor/geofm-finetuning/benchmark-geo-bench-paolo/
-      name: test2
-      log_graph: false
-      default_hp_metric: true
-      prefix: ''
-      comment: ''
-      max_queue: 10
-      flush_secs: 120
-      filename_suffix: ''
+  logger: true
   callbacks:
   - class_path: lightning.pytorch.callbacks.RichProgressBar
     init_args:
@@ -171,24 +160,6 @@ data:
     use_metadata: false
 out_dtype: int16
 deploy_config_file: true
-ModelCheckpoint:
-  filename: '{epoch}'
-  monitor: val/loss
-  verbose: false
-  save_top_k: 1
-  mode: min
-  save_weights_only: false
-  auto_insert_metric_name: true
-  enable_version_counter: true
-StateDictModelCheckpoint:
-  filename: '{epoch}_state_dict'
-  monitor: val/loss
-  verbose: false
-  save_top_k: 1
-  mode: min
-  save_weights_only: true
-  auto_insert_metric_name: true
-  enable_version_counter: true
 optimizer:
   class_path: torch.optim.AdamW
   init_args:

inference.py

@@ -185,7 +185,7 @@ def run_model(input_data, temporal_coords, location_coords, model, datamodule, i
     for x in windows:
         # Apply standardization
         x = datamodule.test_transform(image=x.squeeze().numpy().transpose(1,2,0))
-        x = datamodule.aug(x)['image']
+        x = datamodule.aug(x['image'])
 
         with torch.no_grad():
             x = x.to(model.device)
@@ -317,7 +317,7 @@ if __name__ == "__main__":
     parser.add_argument(
         "--checkpoint",
         type=str,
-        default="Prithvi-EO-V2-300M-TL-Sen1Floods11.ckpt",
+        default="Prithvi-EO-V2-300M-TL-Sen1Floods11.pt",
         help="Path to a checkpoint file to load from.",
     )
     parser.add_argument(