fix: use the correct parameter name in torch.unique call

as in https://pytorch.org/docs/2.1/generated/torch.unique.html#torch-unique
imsoumya18 · Dec 8, 2023 · f462710 · f462710
1 parent dec9ce6
commit f462710
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/ivy/functional/backends/torch/set.py b/ivy/functional/backends/torch/set.py
@@ -114,7 +114,7 @@ def unique_inverse(
         x = torch.flatten(x)
         axis = 0
 
-    values, inverse_indices = torch.unique(x, return_inverse=True, axis=axis)
+    values, inverse_indices = torch.unique(x, return_inverse=True, dim=axis)
     nan_idx = torch.isnan(x)
     if nan_idx.any():
         inverse_indices[nan_idx] = torch.where(torch.isnan(values))[0][0]