Forward script fallback.

Tim-Salzmann · Jan 11, 2024 · ac4f800 · ac4f800
1 parent 1f85dcd
commit ac4f800
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/l4casadi/l4casadi.py b/l4casadi/l4casadi.py
@@ -266,7 +266,9 @@ def export_torch_traces(self, rows: int, cols: int) -> Tuple[bool, bool]:
 
         out_folder = self.build_dir
 
-        torch.jit.trace(self.model, d_inp).save((out_folder / f'{self.name}_forward.pt').as_posix())
+        self._jit_compile_and_save(make_fx(functionalize(self.model, remove='mutations_and_views'))(d_inp),
+                                   (out_folder / f'{self.name}_forward.pt').as_posix(),
+                                   d_inp)
 
         exported_jacrev = False
         if self._with_jacobian: