From e788d5b18e9075cd3d72c6def79b2433d1f9fb52 Mon Sep 17 00:00:00 2001
From: Ali Shibli <30579166+alishibli97@users.noreply.github.com>
Date: Thu, 17 Oct 2024 19:43:21 +0200
Subject: [PATCH] fix checkpoint for ssl4eo_mae_optical (#98)

---
 pangaea/encoders/ssl4eo_mae_encoder.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pangaea/encoders/ssl4eo_mae_encoder.py b/pangaea/encoders/ssl4eo_mae_encoder.py
index 68bbc1b..162158a 100644
--- a/pangaea/encoders/ssl4eo_mae_encoder.py
+++ b/pangaea/encoders/ssl4eo_mae_encoder.py
@@ -158,7 +158,7 @@ def forward(self, image):
 
     def load_encoder_weights(self, logger: Logger) -> None:
         checkpoint = torch.load(self.encoder_weights, map_location="cpu")
-        pretrained_model = checkpoint["model"]
+        pretrained_model = checkpoint["state_dict"]
 
         k = pretrained_model.keys()
         pretrained_encoder = {}