docstr: Performs a train step.
  
  Parameters
  ----------
  x
      Input tensor.
  y
      Label tensor.
docstr_tok: ["Performs","a","train","step",".","Parameters","--","--","--","--","--","x","Input","tensor",".","y","Label","tensor","."]
code: def train_step(
      self, x: torch.Tensor, y: Union[torch.Tensor, List[torch.Tensor]]
  ) -> Dict[str, float]:
      """
      Performs a train step.
  
      Parameters
      ----------
      x
          Input tensor.
      y
          Label tensor.
      """
      # set model to train
      self.train()
  
      # send tensors to device
      x = x.to(self.device)
      y_true: Union[torch.Tensor, List[torch.Tensor]] = (
          [y_i.to(self.device) for y_i in y]
          if isinstance(y, list)
          else y.to(self.device)
      )
  
      # compute output
      y_pred: Union[torch.Tensor, List[torch.Tensor]] = self.forward(x)
  
      # validate prediction and labels
      self.validate_prediction_labels(y_pred=y_pred, y_true=y_true)
  
      # compute loss
      loss, results = self.compute_loss(y_pred=y_pred, y_true=y_true)
  
      # perform gradient update
      self.optimizer.zero_grad()
      loss.backward()
      self.optimizer.step()
  
      # update metrics
      metrics = self.compute_metrics(y_pred=y_pred, y_true=y_true)
      results.update(metrics)
      return results
  
imports: import torch
  import torch.nn
  import torch.optim
  from torch.utils.data import DataLoader
  from tqdm import tqdm
  from typing import List, Dict, Callable, Union, Tuple, Optional
  from alibi.models.pytorch.metrics import Metric, LossContainer
label_desc: Use image node_id 6 for calling the Model obj's underlying member method code with example usage: obj.train_step(x, y) and returns: results