class_docstr: Lightning module which trains an MLP until convergence.
class_docstr_tok: ["Lightning","module","which","trains","an","MLP","until","convergence","."]
code: def training_step(
      self,
      batch: Tuple[torch.Tensor, torch.Tensor, torch.Tensor],
      _batch_idx: int,
  ) -> torch.Tensor:
      X, y_true, group_ids = batch
      y_pred = self.model(X)
  
      if self.uses_ranking:
          loss = self.loss(y_pred, y_true, group_ids)
      else:
          loss = self.loss(y_pred, y_true)
  
      self.log("train_loss", loss, on_step=False, on_epoch=True)
      return loss
  
file_path: gluonts/src/gluonts/nursery/tsbench/src/tsbench/surrogate/torch/mlp_lightning_module.py
imports: from typing import List, Tuple
  import pytorch_lightning
  import torch
  from pytorch_lightning.callbacks import Callback, EarlyStopping
  from torch import nn, optim
  from .losses import ListMLELoss
label_desc: Use image node_id 4 for calling the MLPLightningModule obj's underlying member method code with example usage: obj.training_step(batch, _batch_idx) and returns: loss