From 29a84828f0339094bd9c50e7421dda3b1ba034ed Mon Sep 17 00:00:00 2001
From: Saoirse Stewart <saoirse.stewart@arm.com>
Date: Wed, 11 Sep 2024 15:14:43 +0100
Subject: [PATCH] Update aot_arm_compiler to use export_for_training

Change-Id: I5e92b9574d0950ef2f4da413f319391b2d89c333
---
 examples/arm/aot_arm_compiler.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/examples/arm/aot_arm_compiler.py b/examples/arm/aot_arm_compiler.py
index 9a45195e58f..e68bd8a557d 100644
--- a/examples/arm/aot_arm_compiler.py
+++ b/examples/arm/aot_arm_compiler.py
@@ -196,7 +196,7 @@ def forward(self, x):
     model = model.eval()
 
     # pre-autograd export. eventually this will become torch.export
-    model = torch._export.capture_pre_autograd_graph(model, example_inputs)
+    model = torch.export.export_for_training(model, example_inputs).module()
 
     # Quantize if required
     if args.quantize: