def slots(self, slots):
    """Sets the slots of the optimizers and this class (replicated)."""
    self._optimizer.slots = slots
    self._slots = tl.on_cpu(tl.for_n_devices(slots, self._n_devices))
Sets the slots of the optimizers and this class (replicated).
trax/optimizers/trainer.py