def update_weights_and_state(self, weights=None, state=None):
    """Updates the weights and state of the trained model.

    Sends this data both to the singleton model accessible via Loop.model
    and to the replicated model on the accelerator.

    Useful when the weights or state are modified outside of training, e.g.
    during data collection in RL agents.

    Args:
      weights: Model weights or ``None``. If ``None``, don't set.
      state: Model state or ``None``. If ``None``, don't set.
    """
    for trainer in self._trainer_per_task:
      acc_model_with_loss = trainer.accelerated_model_with_loss
      if weights is not None:
        self._model.weights = weights
        acc_model_with_loss.replicate_weights(trainer.model_with_loss.weights)
      if state is not None:
        self._model.state = state
        acc_model_with_loss.replicate_state(trainer.model_with_loss.state)
Updates the weights and state of the trained model.

    Sends this data both to the singleton model accessible via Loop.model
    and to the replicated model on the accelerator.

    Useful when the weights or state are modified outside of training, e.g.
    during data collection in RL agents.

    Args:
      weights: Model weights or ``None``. If ``None``, don't set.
      state: Model state or ``None``. If ``None``, don't set.
    
trax/supervised/training.py