Use at's return value

Author: NeilGirdhar

src/array_api_extra/_lib/_funcs.py

@@ -420,10 +420,10 @@ def one_hot(
         out = xp.zeros((x.size, num_classes), dtype=dtype)
         x_flattened = xp.reshape(x, (-1,))
         if is_numpy_namespace(xp):
-            at(out)[xp.arange(x_size), x_flattened].set(1)
+            out = at(out)[xp.arange(x_size), x_flattened].set(1)
         else:
             for i in range(x_size):
-                at(out)[i, int(x_flattened[i])].set(1)
+                out = at(out)[i, int(x_flattened[i])].set(1)
     if x.ndim != 1:
         out = xp.reshape(out, (*x.shape, num_classes))
     if axis != -1: