diff --git a/onnxscript/function_libs/torch_lib/ops/core.py b/onnxscript/function_libs/torch_lib/ops/core.py index 99c20eca18..48c64df4d4 100644 --- a/onnxscript/function_libs/torch_lib/ops/core.py +++ b/onnxscript/function_libs/torch_lib/ops/core.py @@ -3246,7 +3246,10 @@ def aten_div_mode(self: TReal, other: TReal, rounding_mode: Optional[str] = None if rounding_mode == "trunc": # Rounds the results of the division towards zero. # Equivalent to C-style integer division - return aten_trunc(op.Div(self, other)) + quotient = op.Div(self, other) + if self.dtype == ir.DataType.FLOAT16: + quotient = op.Cast(quotient, to=FLOAT16.dtype) + return aten_trunc(quotient) if rounding_mode == "floor": return op.Floor(op.Div(self, other)) diff --git a/tests/function_libs/torch_lib/ops_test_data.py b/tests/function_libs/torch_lib/ops_test_data.py index 3658a35b9d..e6c187fda2 100644 --- a/tests/function_libs/torch_lib/ops_test_data.py +++ b/tests/function_libs/torch_lib/ops_test_data.py @@ -676,16 +676,9 @@ def _where_input_wrangler( ), TorchLibOpInfo("true_divide", core_ops.aten_div), TorchLibOpInfo("true_divide", core_ops.aten_div_complex, complex=True), - TorchLibOpInfo("div_mode", core_ops.aten_div_mode) - .skip( + TorchLibOpInfo("div_mode", core_ops.aten_div_mode).skip( variant_name="no_rounding_mode", reason="this variation requires the rounding_mode argument", - ) - .skip( - variant_name="trunc_rounding", - dtypes=(torch.float16,), - # Numbers match sometimes but not other times - reason="fixme: off-by-one. https://github.com/microsoft/onnxscript/issues/990", ), TorchLibOpInfo("dot", core_ops.aten_dot), TorchLibOpInfo(