def METHOD_NAME(device):
    metric_devices = [torch.device("cpu")]
    if device.type != "xla":
        metric_devices.append(idist.device())
    for metric_device in metric_devices:
        mpd = MeanPairwiseDistance(device=metric_device)
        for dev in [mpd._device, mpd._sum_of_distances.device]:
            assert dev == metric_device, f"{type(dev)}:{dev} vs {type(metric_device)}:{metric_device}"
        y_pred = torch.tensor([[3.0, 4.0], [-3.0, -4.0]])
        y = torch.zeros(2, 2)
        mpd.update((y_pred, y))
        for dev in [mpd._device, mpd._sum_of_distances.device]:
            assert dev == metric_device, f"{type(dev)}:{dev} vs {type(metric_device)}:{metric_device}"