[FIX] AttributeError: 'ExperimentWriter' object has no attribute 'add_figure'

pytorch_forecasting/models/base_model.py

@@ -976,6 +976,10 @@ def log_prediction(
             if not mpl_available:
                 return None  # don't log matplotlib plots if not available
 
+            # Don't log figures if add_figure is not available
+            if not hasattr(self.logger.experiment, "add_figure"):
+                return None
+
             for idx in log_indices:
                 fig = self.plot_prediction(x, out, idx=idx, add_loss_to_title=True, **kwargs)
                 tag = f"{self.current_stage} prediction"
@@ -1149,6 +1153,10 @@ def log_gradient_flow(self, named_parameters: Dict[str, torch.Tensor]) -> None:
         if not mpl_available:
             return None
 
+        # Don't log figures if add_figure is not available
+        if not hasattr(self.logger.experiment, "add_figure"):
+            return None
+
         import matplotlib.pyplot as plt
 
         fig, ax = plt.subplots()

pytorch_forecasting/models/nbeats/__init__.py

@@ -268,6 +268,10 @@ def log_interpretation(self, x, out, batch_idx):
         if not mpl_available:
             return None
 
+        # Don't log figures if add_figure is not available
+        if not hasattr(self.logger.experiment, "add_figure"):
+            return None
+
         label = ["val", "train"][self.training]
         if self.log_interval > 0 and batch_idx % self.log_interval == 0:
             fig = self.plot_interpretation(x, out, idx=0)

pytorch_forecasting/models/nhits/__init__.py

@@ -544,6 +544,10 @@ def log_interpretation(self, x, out, batch_idx):
         if not mpl_available:
             return None
 
+        # Don't log figures if add_figure is not available
+        if not hasattr(self.logger.experiment, "add_figure"):
+            return None
+
         label = ["val", "train"][self.training]
         if self.log_interval > 0 and batch_idx % self.log_interval == 0:
             fig = self.plot_interpretation(x, out, idx=0)

pytorch_forecasting/models/temporal_fusion_transformer/__init__.py

@@ -821,6 +821,10 @@ def log_interpretation(self, outputs):
         if not mpl_available:
             return None
 
+        # Don't log figures if add_figure is not available
+        if not hasattr(self.logger.experiment, "add_figure"):
+            return None
+
         import matplotlib.pyplot as plt
 
         figs = self.plot_interpretation(interpretation)  # make interpretation figures
@@ -857,6 +861,11 @@ def log_embeddings(self):
         """
         Log embeddings to tensorboard
         """
+
+        # Don't log embeddings if add_embedding is not available
+        if not hasattr(self.logger.experiment, "add_embedding"):
+            return None
+
         for name, emb in self.input_embeddings.items():
             labels = self.hparams.embedding_labels[name]
             self.logger.experiment.add_embedding(