[torchlib] Fix aten_mean_dim (#1962)

Fix when the rank of `dim` is not known, the conditional will pick the wrong branch because the truthiness of a value is always True.
microsoft · Nov 22, 2024 · 99b3d26 · 99b3d26
1 parent 5a4d22e
commit 99b3d26
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/onnxscript/function_libs/torch_lib/ops/core.py b/onnxscript/function_libs/torch_lib/ops/core.py
@@ -5225,9 +5225,8 @@ def aten_mean_dim(self: TReal, dim: INT64, keepdim: bool = False) -> TReal:
     if IsScalar(self):
         result = self
     else:
-        if IsScalar(dim):
-            dim = op.Unsqueeze(dim, axes=0)
-        result = op.ReduceMean(self, dim, keepdims=keepdim)
+        dims = op.Reshape(dim, op.Constant(value_ints=[-1]))
+        result = op.ReduceMean(self, dims, keepdims=keepdim)
     return result