def override_get_stream():
    r"""
    This function overrides `_get_stream` in PyTorch to provide XPU support.
    """

    def _get_stream(device: int):
        r"""
        Gets a background stream for copying between CPU and XPU.
        """
        global _streams
        if device == -1:
            return None
        if _streams is None:
            _streams = [None] * torch.xpu.device_count()
        if _streams[device] is None:
            _streams[device] = torch.xpu.Stream(device)
        return _streams[device]

    torch.nn.parallel._functions._get_stream = _get_stream
    return _get_stream

    This function overrides `_get_stream` in PyTorch to provide XPU support.
    
intel_extension_for_pytorch/xpu/overrides.py