Fix BatchKey

openclimatefix · Jan 8, 2024 · 6142796 · 6142796
1 parent d171ab0
commit 6142796
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/pvnet/models/multimodal/multimodal.py b/pvnet/models/multimodal/multimodal.py
@@ -234,7 +234,7 @@ def forward(self, x):
             sat_data = x[BatchKey.satellite_actual][:, : self.sat_sequence_len]
             sat_data = torch.swapaxes(sat_data, 1, 2).float()  # switch time and channels
             if self.add_image_embedding_channel:
-                id = x[BatchKey.sensor_id][:, 0].int()
+                id = x[BatchKey.wind_id][:, 0].int()
                 sat_data = self.sat_embed(sat_data, id)
             modes["sat"] = self.sat_encoder(sat_data)
 
@@ -264,7 +264,7 @@ def forward(self, x):
 
         # ********************** Embedding of GSP ID ********************
         if self.embedding_dim:
-            id = x[BatchKey.sensor_id][:, 0].int()
+            id = x[BatchKey.wind_id][:, 0].int()
             id_embedding = self.embed(id)
             modes["id"] = id_embedding