class_docstr: Callback triggered before and after a training step.
docstr: Called by Loop after training steps, when call_at returned True.
docstr_tok: ["Called","by","Loop","after","training","steps",",","when","call_at","returned","True","."]
code: def on_step_end(self, step):
      """Called by Loop after training steps, when call_at returned True."""
      raise NotImplementedError
  
imports: import collections
  import os
  import gin
  import numpy
  from trax import jaxboard
  from trax import layers
  from trax import shapes
  from trax.rl import serialization_utils
  from trax.supervised import decoding
label_desc: Use image node_id 4 for calling the TrainingStepCallback obj's underlying member method code with example usage: obj.on_step_end(step) without return types