code: def reset_states(self):
      for key in self._states:
          self._states[key]["sum"].assign(
              tf.zeros(
                  tf.shape(self._states[key]["sum"]), dtype=tf.float32
              )
          )
          self._states[key]["count"].assign(0.0)
  
imports: import tensorflow
  import numpy
label_desc: Use image node_id 2 for calling the DictMean obj's underlying member method code with example usage: obj.reset_states() without return types