Move a rollout to a selected device

def to_device(self, device, non_blocking=True):
        """ Move a rollout to a selected device """
        return Transitions(
            size=self.size,
            environment_information=self.environment_information,
            transition_tensors={k: v.to(device, non_blocking=non_blocking) for k, v in self.transition_tensors.items()},
            extra_data=self.extra_data
        )