Fix casting in lerp

microsoft · Jul 16, 2024 · 246efd8 · 246efd8
1 parent 36f96d5
commit 246efd8
Showing 1 changed file with 5 additions and 4 deletions.
diff --git a/onnxscript/function_libs/torch_lib/ops/core.py b/onnxscript/function_libs/torch_lib/ops/core.py
@@ -4663,14 +4663,15 @@ def aten_le_bool(self: BOOL, other: BOOL) -> BOOL:
 
 
 @torch_op(("aten::lerp.Tensor", "aten::lerp.Scalar"))
-def aten_lerp(self: TReal, end: TReal, weight: TReal) -> TReal:
+def aten_lerp(self: TTensor, end: TTensor, weight: TTensor) -> TTensor:
     """lerp.Tensor(Tensor self, Tensor end, Tensor weight) -> Tensor"""
 
-    diff = op.CastLike(op.Sub(end, self), weight)
+    weight = op.CastLike(weight, self)
+    diff = op.Sub(end, self)
     return op.Where(
         op.Less(weight, 0.5),
-        op.Add(self, op.CastLike(op.Mul(weight, diff), self)),
-        op.Sub(end, op.CastLike(op.Mul(diff, op.Sub(1.0, weight)), end))
+        op.Add(self, op.Mul(weight, diff)),
+        op.Sub(end, op.Mul(diff, op.Sub(1.0, weight)))
     )