class_docstr: Class with all code to run experiment.
class_docstr_tok: ["Class","with","all","code","to","run","experiment","."]
code: def train_op(self, images, labels, cur_epoch):
      cur_epoch = cur_epoch[0]  # because cur_epoch is array of size 1
      grads, (_, losses_dict) = self.compute_grads_loss(images, labels)
      grads = objax.functional.parallel.pmean(grads)
      losses_dict = objax.functional.parallel.pmean(losses_dict)
      learning_rate = self.learning_rate(cur_epoch)
      self.optimizer(learning_rate, grads)
      return dict(
          **losses_dict, learning_rate=learning_rate, epoch=cur_epoch
      )
  
file_path: objax/examples/image_classification/imagenet_resnet50_train.py
imports: import os
  import time
  import jax
  import jax.numpy
  import numpy
  from absl import app, flags, logging
  import objax
  from examples.image_classification import imagenet_resnet50_data
  from objax.zoo.resnet_v2 import ResNet50
label_desc: Use image node_id 6 for calling the Experiment obj's underlying member method code with example usage: obj.train_op(images, labels, cur_epoch) and returns: dict