def _revert_to_original_device(self, tensor_distributed, orig_data):
        """
        Internal utility used to move the tensor back to the original device
        (if needed).

        :param: The tensor obtained from a torch.distributed API call.
        :param: The descriptor in the format of
            :meth:`_prepare_for_distributed_comm`.
        :return: The tensor moved to the appropriate device.
        """

        original_device, copy_back, tensor = orig_data
        if copy_back:
            if tensor is None:
                tensor = tensor_distributed.to(original_device)
            else:
                tensor[:] = tensor_distributed

        return tensor

        Internal utility used to move the tensor back to the original device
        (if needed).

        :param: The tensor obtained from a torch.distributed API call.
        :param: The descriptor in the format of
            :meth:`_prepare_for_distributed_comm`.
        :return: The tensor moved to the appropriate device.
        
avalanche/distributed/distributed_helper.py