code: def wrapped_tf_function(points, label):
      """Performs one step of minimization of the loss."""
      # --- subsampling (order DO matter)
      points = points[0 : FLAGS.num_points, ...]
  
      # --- augmentation
      if FLAGS.augment:
          points = tf.map_fn(augment.rotate, points)
          points = augment.jitter(points)
  
      # --- training
      with tf.GradientTape() as tape:
          logits = model(points, training=True)
          loss = model.loss(label, logits)
      variables = model.trainable_variables
      gradients = tape.gradient(loss, variables)
      optimizer.apply_gradients(zip(gradients, variables))
      return loss
  
file_path: graphics/tensorflow_graphics/projects/pointnet/train.py
imports: import tensorflow
  from tensorflow_graphics.datasets import modelnet40
  from tensorflow_graphics.nn.layer import pointnet
  import tqdm
  from .None import augment
  from .None import helpers
label_desc: Use image node_id 1 for calling a global function with example usage: wrapped_tf_function(points, label) and returns: loss