Add allgather check for xpu (#2199)

* add  allgather check for xpu

* style fix

* fix test

* fix test and review

src/accelerate/utils/operations.py

@@ -298,13 +298,6 @@ def _gpu_gather_one(tensor):
         if not tensor.is_contiguous():
             tensor = tensor.contiguous()
 
-        # Check if `tensor` is not on CUDA
-        if state.device.type == "cuda" and tensor.device.type != "cuda":
-            raise RuntimeError(
-                "One or more of the tensors passed to `gather` were not on the GPU while the `Accelerator` is configured for CUDA. "
-                "Please move it to the GPU before calling `gather`."
-            )
-
         if state.backend is not None and state.backend != "gloo":
             # We use `empty` as `all_gather_into_tensor` slightly
             # differs from `all_gather` for better efficiency,
@@ -351,6 +344,11 @@ def wrapper(*args, **kwargs):
             tensor = kwargs["tensor"]
         else:
             tensor = args[0]
+        if PartialState().device.type != find_device(tensor).type:
+            raise DistributedOperationException(
+                f"One or more of the tensors passed to {operation} were not on the {tensor.device.type} while the `Accelerator` is configured for {PartialState().device.type}. "
+                f"Please move it to the {PartialState().device.type} before calling {operation}."
+            )
         shapes = get_shape(tensor)
         output = gather_object([shapes])
         if output[0] is not None: