Skip to content

Commit 9f82a6c

Browse files
committed
enable aoti for preprocess ci
ghstack-source-id: e1f8ab2 Pull Request resolved: #6553
1 parent 6fd92ab commit 9f82a6c

File tree

1 file changed

+10
-10
lines changed

1 file changed

+10
-10
lines changed

examples/models/llama3_2_vision/preprocess/test_preprocess.py

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -76,10 +76,10 @@ def initialize_models(resize_to_max_canvas: bool) -> Dict[str, Any]:
7676
strict=False,
7777
)
7878

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+
)
8383

8484
edge_program = to_edge(
8585
exported_model, compile_config=EdgeCompileConfig(_check_ir_validity=False)
@@ -91,7 +91,7 @@ def initialize_models(resize_to_max_canvas: bool) -> Dict[str, Any]:
9191
"reference_model": reference_model,
9292
"model": model,
9393
"exported_model": exported_model,
94-
# "aoti_path": aoti_path,
94+
"aoti_path": aoti_path,
9595
"executorch_model": executorch_model,
9696
}
9797

@@ -237,11 +237,11 @@ def run_preprocess(
237237
self.assertEqual(reference_ar, et_ar.tolist())
238238

239239
# 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())
245245

246246
# This test setup mirrors the one in torchtune:
247247
# https://github.com/pytorch/torchtune/blob/main/tests/torchtune/models/clip/test_clip_image_transform.py

0 commit comments

Comments
 (0)