class_docstr: A callback that stores plots of the  predictions for a collection of samples every n epochs.
  The plots display the query (past and future), forecasted quantiles and the time series in the support set
  of this sample aligned with their attention scores. This callback works only for models with attention mechanism!
  
  Args:
      log_batch: For each sample in the batch the prediction is plotted when the callback is called.
          The batch size should not be too large (i.e. < 10).
      quantiles: The quantiles that are predicted.
      split: Specifies the split the batch comes from (i.e. from training or validation split)
      every_n_epochs: Specifies how often the plots are generated.
          Setting this to a large value can save time since plotting can be time consuming
          (especially when small datasets are used).
class_docstr_tok: ["A","callback","that","stores","plots","of","the","predictions","for","a","collection","of","samples","every","n","epochs",".","The","plots","display","the","query","(","past","and","future",")",",","forecasted","quantiles","and","the","time","series","in","the","support","set","of","this","sample","aligned","with","their","attention","scores",".","This","callback","works","only","for","models","with","attention","mechanism","!","Args",":","log_batch",":","For","each","sample","in","the","batch","the","prediction","is","plotted","when","the","callback","is","called",".","The","batch","size","should","not","be","too","large","(","i.e",".","<","10",")",".","quantiles",":","The","quantiles","that","are","predicted",".","split",":","Specifies","the","split","the","batch","comes","from","(","i.e",".","from","training","or","validation","split",")","every_n_epochs",":","Specifies","how","often","the","plots","are","generated",".","Setting","this","to","a","large","value","can","save","time","since","plotting","can","be","time","consuming","(","especially","when","small","datasets","are","used",")","."]
code: def on_validation_epoch_end(
      self, trainer: Trainer, pl_module: MetaLightningModule
  ) -> None:
      if trainer.current_epoch % self.every_n_epochs:
          return
      # Bring the tensors to CPU
      log_batch = self.log_batch.to(device=pl_module.device)
      query_past = log_batch.query_past
      query_future = log_batch.query_future
      support_set = log_batch.support_set
  
      # Get model prediction
      pred, attention = pl_module.model(
          supps=support_set,
          query=query_past,
          return_attention=True,
      )
      pred = SeriesBatch(
          pred, query_future.lengths, query_future.split_sections
      ).unpad(to_numpy=True, squeeze=True)
      qp = query_past.unpad(to_numpy=True, squeeze=True)
      qf = query_future.unpad(to_numpy=True, squeeze=True)
      supps = support_set.unpad(to_numpy=True, squeeze=False)
      n_supps = support_set.split_sections[0]
      attention = attention.reshape(
          n_supps * attention.size()[0], 1, -1
      )
      attention = SeriesBatch(
          attention.transpose(1, 2),
          support_set.lengths,
          support_set.split_sections,
      ).unpad(to_numpy=True, squeeze=False)
      # attention = [[tensor_to_np(att) for att in supps] for supps in attention]
  
      fig = plot_forecast_supportset_attention(
          qp,
          qf,
          pred,
          supps=supps,
          attention=attention,
          quantiles=self.quantiles,
      )
      fig.savefig(
          get_save_dir_from_csvlogger(trainer.logger)
          / f"pred_supps_{self.split + '_' if self.split else ''}ep{trainer.current_epoch}.png",
          bbox_inches="tight",
      )
      plt.close(fig)
  
imports: from typing import Optional, Union, List
  from pytorch_lightning import Callback, Trainer
  import matplotlib.pyplot
  from meta.data.batch import TripletBatch
  from meta.models.module import MetaLightningModule
  from .common import get_save_dir_from_csvlogger, get_loss_steps
  from meta.data.batch import SeriesBatch
  from meta.vis.forecast import plot_quantile_forecast, plot_forecast_supportset_attention
label_desc: Use image node_id 2 for calling the ForecastSupportSetAttentionPlotLoggerCallback obj's underlying member method code with example usage: obj.on_validation_epoch_end(trainer, pl_module) without return types