Make wandb work again with pytorch_lightning.logger

neural_lam/models/ar_model.py

@@ -568,8 +568,16 @@ def aggregate_and_plot_metrics(self, metrics_dict, prefix):
                     )
                 )
 
+        # Ensure that log_dict has structure for logging as dict(str, plt.Figure)
+        assert all(
+            isinstance(key, str) and isinstance(value, plt.Figure)
+            for key, value in log_dict.items()
+        )
+
         if self.trainer.is_global_zero and not self.trainer.sanity_checking:
-            self.logger.log_image(log_dict)  # Log all
+            for key, figure in log_dict.items():
+                self.logger.log_image(key=key, images=[figure])
+
             plt.close("all")  # Close all figs
 
     def on_test_epoch_end(self):