code: def train_step(dense_features, sparse_features, label):
      with tf.GradientTape() as tape:
          loss_value = dlrm_model(
              dense_features, sparse_features, label
          )
      gradients = tape.gradient(
          loss_value, dlrm_model.trainable_variables
      )
      optimizer.apply_gradients(
          zip(gradients, dlrm_model.trainable_variables)
      )
      return loss_value
  
file_path: openrec/tf2_examples/dlrm_criteo.py
imports: from tensorflow.data import Dataset
  from openrec.tf2.recommenders import DLRM
  from tensorflow.keras import optimizers
  from tqdm import tqdm
  import tensorflow
  import dataloader
label_desc: Use image node_id 1 for calling a global function with example usage: train_step(dense_features, sparse_features, label) and returns: loss_value