def _log_n_weights(self):
    """"Logs the number of weights in the training model."""
    def _size(x):
      try:
        return x.size
      except Exception:  # pylint: disable=broad-except
        return 0
    sizes = fastmath.nested_map(_size, self._model.weights)
    total_size = sum(fastmath.tree_flatten(sizes))
    total_size *= base.N_WEIGHTS_SHARDS
    self._log_step('Total number of trainable weights: %d' % total_size)
"Logs the number of weights in the training model.
trax/supervised/training.py