class_docstr: Callback triggered before and after a training step.
docstr: Returns whether the callback should be called at a given step.
docstr_tok: ["Returns","whether","the","callback","should","be","called","at","a","given","step","."]
code: def call_at(self, step):
      """Returns whether the callback should be called at a given step."""
      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 2 for calling the TrainingStepCallback obj's underlying member method code with example usage: obj.call_at(step) without return types