fix: array squeeze in torch frontend median

ivy-llc · Sep 10, 2024 · c686786 · c686786
1 parent 71af269
commit c686786
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/ivy/functional/frontends/torch/reduction_ops.py b/ivy/functional/frontends/torch/reduction_ops.py
@@ -172,9 +172,9 @@ def median(input, dim=None, keepdim=False, *, out=None):
         median_indices = ivy.gather(
             sorted_indices, (sorted_indices.shape[dim] - 1) // 2, axis=dim
         )
-        median_values = ivy.take_along_axis(
+        median_values = ivy.squeeze(ivy.take_along_axis(
             input, ivy.expand_dims(median_indices, axis=dim), dim
-        ).squeeze(axis=dim)
+        ), axis=dim)
 
         if keepdim:
             median_values = ivy.expand_dims(median_values, axis=dim)