docstr: Override `training_step` from PredictionAlgorithm class for ERM-specific training loop.
docstr_tok: ["Override","`","training_step","`","from","PredictionAlgorithm","class","for","ERM-specific","training","loop","."]
code: def training_step(self, train_batch, batch_idx):
      """
      Override `training_step` from PredictionAlgorithm class for ERM-specific training loop.
  
      """
  
      x = torch.cat([x for x, y, _ in train_batch])
      y = torch.cat([y for x, y, _ in train_batch])
  
      out = self.model(x)
      loss = F.cross_entropy(out, y)
      acc = (torch.argmax(out, dim=1) == y).float().mean()
  
      metrics = {"train_acc": acc, "train_loss": loss}
      self.log_dict(
          metrics, on_step=False, on_epoch=True, prog_bar=True
      )
  
      return loss
  
imports: import torch
  from torch.nn import functional
  from dowhy.causal_prediction.algorithms.base_algorithm import PredictionAlgorithm
label_desc: Use image node_id 2 for calling the ERM obj's underlying member method code with example usage: obj.training_step(train_batch, batch_idx) and returns: loss