def METHOD_NAME(self, gpus, chunk_sizes, device):
    if len(gpus) > 1:
        self.model_with_loss = DataParallel(
            self.model_with_loss, device_ids=gpus, chunk_sizes=chunk_sizes
        ).to(device)
    else:
        self.model_with_loss = self.model_with_loss.to(device)
    for state in self.optimizer.state.values():
        for k, v in state.items():
            if isinstance(v, torch.Tensor):
                state[k] = v.to(device=device, non_blocking=True)