Update src/array_api_extra/_lib/_funcs.py

Co-authored-by: Guido Imperiale <[email protected]>

Author: NeilGirdhar

src/array_api_extra/_lib/_funcs.py

@@ -394,7 +394,7 @@ def one_hot(
         # specification.
         msg = "x must have a concrete size."
         raise TypeError(msg)
-    out = xp.zeros((x.size, num_classes), dtype=dtype)
+    out = xp.zeros((x.size, num_classes), dtype=dtype, device=_compat.device(x))
     x_flattened = xp.reshape(x, (-1,))
     if supports_fancy_indexing:
         out = at(out)[xp.arange(x_size), x_flattened].set(1)