Prithvi add loss dict

Signed-off-by: Benedikt Blumenstiel <[email protected]>
IBM · Feb 21, 2025 · f6e6b70 · f6e6b70
1 parent e282b09
commit f6e6b70
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/terratorch/models/backbones/prithvi_mae.py b/terratorch/models/backbones/prithvi_mae.py
@@ -755,6 +755,7 @@ def forward(
         latent, mask, ids_restore = self.encoder(pixel_values, temporal_coords, location_coords, mask_ratio)
         pred = self.decoder(latent, ids_restore, temporal_coords, location_coords, input_size=pixel_values.shape)
         loss = self.forward_loss(pixel_values, pred, mask)
+        loss = {'loss': loss}  # TerraTorch expects loss dict
         return loss, pred, mask
 
     def forward_features(