@@ -76,10 +76,10 @@ def initialize_models(resize_to_max_canvas: bool) -> Dict[str, Any]:
76
76
strict = False ,
77
77
)
78
78
79
- # aoti_path = torch._inductor.aot_compile(
80
- # exported_model.module(),
81
- # model.get_example_inputs(),
82
- # )
79
+ aoti_path = torch ._inductor .aot_compile (
80
+ exported_model .module (),
81
+ model .get_example_inputs (),
82
+ )
83
83
84
84
edge_program = to_edge (
85
85
exported_model , compile_config = EdgeCompileConfig (_check_ir_validity = False )
@@ -91,7 +91,7 @@ def initialize_models(resize_to_max_canvas: bool) -> Dict[str, Any]:
91
91
"reference_model" : reference_model ,
92
92
"model" : model ,
93
93
"exported_model" : exported_model ,
94
- # "aoti_path": aoti_path,
94
+ "aoti_path" : aoti_path ,
95
95
"executorch_model" : executorch_model ,
96
96
}
97
97
@@ -237,11 +237,11 @@ def run_preprocess(
237
237
self .assertEqual (reference_ar , et_ar .tolist ())
238
238
239
239
# Run aoti model and check it matches reference model.
240
- # aoti_path = models["aoti_path"]
241
- # aoti_model = torch._export.aot_load(aoti_path, "cpu")
242
- # aoti_image, aoti_ar = aoti_model(image_tensor, inscribed_size, best_resolution)
243
- # self.assertTrue(torch.allclose(reference_image, aoti_image))
244
- # self.assertEqual(reference_ar, aoti_ar.tolist())
240
+ aoti_path = models ["aoti_path" ]
241
+ aoti_model = torch ._export .aot_load (aoti_path , "cpu" )
242
+ aoti_image , aoti_ar = aoti_model (image_tensor , inscribed_size , best_resolution )
243
+ self .assertTrue (torch .allclose (reference_image , aoti_image ))
244
+ self .assertEqual (reference_ar , aoti_ar .tolist ())
245
245
246
246
# This test setup mirrors the one in torchtune:
247
247
# https://github.com/pytorch/torchtune/blob/main/tests/torchtune/models/clip/test_clip_image_transform.py
0 commit comments