"""Stage Accelerate's cross-GPU copies through host RAM on this workstation.

Direct CUDA peer copies report success but fail equality checks here. The existing
model-manager profiles also disable NCCL P2P. This adapter changes no driver or
global machine settings; it affects only this Python process.
"""
import torch


def install():
    import accelerate.hooks
    import accelerate.utils
    import accelerate.utils.operations as operations
    if getattr(operations.send_to_device, "_host_staged", False):
        return
    original = operations.send_to_device

    def send_to_device(tensor, device, non_blocking=False, skip_keys=None):
        if isinstance(tensor, torch.Tensor) and tensor.is_cuda:
            destination = torch.device(f"cuda:{device}" if isinstance(device, int) else device)
            if destination.type == "cuda":
                index = torch.cuda.current_device() if destination.index is None else destination.index
                if tensor.device.index != index:
                    return tensor.to("cpu", non_blocking=False).to(destination, non_blocking=False)
        return original(tensor, device, non_blocking=non_blocking, skip_keys=skip_keys)

    send_to_device._host_staged = True
    operations.send_to_device = send_to_device
    accelerate.hooks.send_to_device = send_to_device
    accelerate.utils.send_to_device = send_to_device