Skip to content

Commit 433b751

Browse files
justinchubybmehta001
authored andcommitted
Fix DORT CI according to torch-nightly changes (microsoft#2125)
1 parent f90c2d8 commit 433b751

File tree

2 files changed

+5
-8
lines changed

2 files changed

+5
-8
lines changed

onnxscript/tools/training_helper.py

Lines changed: 4 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -3,16 +3,13 @@
33
# Licensed under the MIT License.
44
# --------------------------------------------------------------------------
55
import torch
6-
from torch.onnx import ExportOptions
7-
from torch.onnx import _OrtBackend as OrtBackend
8-
from torch.onnx import _OrtBackendOptions as OrtBackendOptions
6+
from torch.onnx import _OrtBackend, _OrtBackendOptions
97

108

11-
def make_aot_ort(dynamic: bool = False):
9+
def make_aot_ort():
1210
"""Implements an autograd backend for torch.compile based on onnxrt backend."""
13-
export_options = ExportOptions(dynamic_shapes=dynamic)
14-
options = OrtBackendOptions(export_options=export_options)
15-
ort_backend = OrtBackend(options=options)
11+
options = _OrtBackendOptions()
12+
ort_backend = _OrtBackend(options=options)
1613
return ort_backend
1714

1815

onnxscript/tools/transformers_models/phi_test.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -89,7 +89,7 @@ def test_phi_dort_static(self):
8989
input_tensors = input_tensors_many[0]
9090
expected = model(*input_tensors)
9191

92-
local_aot_ort = onnxscript.tools.training_helper.make_aot_ort(dynamic=False)
92+
local_aot_ort = onnxscript.tools.training_helper.make_aot_ort()
9393

9494
compiled_model = torch.compile(
9595
copy.deepcopy(model),

0 commit comments

Comments
 (0)