remove timestep restriction

Prithvi_run_inference.py

@@ -252,7 +252,7 @@ def main(data_files: List[str], yaml_file_path: str, checkpoint: str, output_dir
         params = yaml.safe_load(f)
 
     # data related
-    num_frames = params['num_frames']
+    num_frames = len(data_files)
     img_size = params['img_size']
     bands = params['bands']
     mean = params['data_mean']
@@ -272,8 +272,9 @@ def main(data_files: List[str], yaml_file_path: str, checkpoint: str, output_dir
 
     mask_ratio = params['mask_ratio'] if mask_ratio is None else mask_ratio
 
-    # We must have *num_frames* files to build one example!
-    assert len(data_files) == num_frames, "File list must be equal to expected number of frames."
+    print(f"\nTreating {len(data_files)} files as {len(data_files)} time steps from the same location\n")
+    if len(data_files) != 3:
+        print("The original model was trained for 3 time steps (expecting 3 files). \nResults with different numbers of timesteps may vary")
 
     if torch.cuda.is_available():
         device = torch.device('cuda')
@@ -310,7 +311,10 @@ def main(data_files: List[str], yaml_file_path: str, checkpoint: str, output_dir
     model.to(device)
 
     state_dict = torch.load(checkpoint, map_location=device)
-    model.load_state_dict(state_dict)
+    # discard fixed pos_embedding weight
+    del state_dict['pos_embed']
+    del state_dict['decoder_pos_embed']
+    model.load_state_dict(state_dict, strict=False)
     print(f"Loaded checkpoint from {checkpoint}")
 
     # Running model --------------------------------------------------------------------------------