class_docstr: PyTorch Lightning module for the meta models.
class_docstr_tok: ["PyTorch","Lightning","module","for","the","meta","models","."]
code: def test_step(
      self, batch: TripletBatch, batch_idx, *args: Any
  ) -> None:
      # a hack since pl considers a list of only one dataloader not as multiple
      if len(self.test_dataset_names) == 1:
          dataset_idx = 0
      else:
          dataset_idx = args[0]
      y_true = batch.query_future.sequences
      mask = get_mask(y_true.size()[1], batch.query_future.lengths)
      y_pred = self.model(batch.support_set, batch.query_past)
      d_name = self.test_dataset_names[dataset_idx]
      self.test_crps[dataset_idx](
          y_pred,
          batch.query_future.sequences,
          mask,
          batch.query_future.scales,
      )
      self.log(
          f"{d_name}_test_crps",
          self.test_crps[dataset_idx],
          add_dataloader_idx=False,
      )
      if self.test_crps_scaled is not None:
          self.test_crps_scaled[dataset_idx](
              y_pred, batch.query_future.sequences, mask
          )
          self.log(
              f"{d_name}_test_crps_scaled",
              self.test_crps_scaled[dataset_idx],
              add_dataloader_idx=False,
          )
      if self.quantile_width is not None:
          self.test_quantile_width[dataset_idx](
              y_pred, batch.query_future.sequences, mask
          )
          self.log(
              f"{d_name}_test_quantile_width",
              self.test_quantile_width[dataset_idx],
              add_dataloader_idx=False,
          )
  
file_path: gluonts/src/gluonts/nursery/few_shot_prediction/src/meta/models/module.py
imports: from typing import Any, Dict, Optional, List
  import pytorch_lightning
  import torch
  import torch.optim
  from torch import nn
  from torchmetrics import Metric
  from torch.optim.lr_scheduler import ReduceLROnPlateau
  from meta.data.batch import TripletBatch
  from meta.models.model import MetaModel
  from meta.common.torch import get_mask
label_desc: Use image node_id 6 for calling the MetaLightningModule obj's underlying member method code with example usage: obj.test_step(batch, batch_idx) without return types