code: def train_step(self, data):
      x, y = data
  
      # Call torch.nn.Module.zero_grad() to clear the leftover gradients
      # for the weights from the previous train step.
      self.zero_grad()
  
      # Compute loss
      y_pred = self(x, training=True)  # Forward pass
      loss = self.loss_fn(y, y_pred)
  
      # Call torch.Tensor.backward() on the loss to compute gradients
      # for the weights.
      loss.backward()
  
      trainable_weights = [v for v in self.trainable_weights]
      gradients = [v.value.grad for v in trainable_weights]
  
      # Update weights
      with torch.no_grad():
          self.optimizer.apply(gradients, trainable_weights)
  
      # Compute our own metrics
      self.loss_tracker.update_state(loss)
      self.mae_metric.update_state(y, y_pred)
      return {
          "loss": self.loss_tracker.result(),
          "mae": self.mae_metric.result(),
      }
  
imports: import os
  import torch
  import keras
  from keras import layers
  import numpy
label_desc: Use image node_id 2 for calling the CustomModel obj's underlying member method code with example usage: obj.train_step(data) and returns: dict