code: def train_step(
      self,
      summary: objax.jaxboard.Summary,
      data: dict,
      step: np.ndarray,
  ):
      kv = self.train_op(step, data["image"], data["label"])
      for k, v in kv.items():
          if jn.isnan(v):
              raise ValueError("NaN", k)
          summary.scalar(k, float(v))
  
imports: import os
  import pickle
  from typing import Optional, Dict, Iterable, Callable
  import jax
  import jax.numpy
  import numpy
  from absl import flags
  from tqdm import tqdm, trange
  import objax
  from examples.fixmatch.libml.data.core import DataSet
  from objax.util import EasyDict
label_desc: Use image node_id 4 for calling the TrainLoopFSL obj's underlying member method code with example usage: obj.train_step(summary, data, step) without return types