Fix source for export to executorch tutorial (#2400)

Summary:
Pull Request resolved: #2400

imported-using-ghimport

Test Plan: Imported from OSS

Reviewed By: angelayi

Differential Revision: D54862906

Pulled By: guangy10

fbshipit-source-id: c26f03f56d58697f155b8201a0711da7dd5b86da

Author: guangy10

docs/source/tutorials_source/export-to-executorch-tutorial.py

@@ -130,11 +130,11 @@ def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
 aten_dialect: ExportedProgram = export(f, example_args)
 
 # Works correctly
-print(aten_dialect(torch.ones(3, 3), torch.ones(3, 3)))
+print(aten_dialect.module()(torch.ones(3, 3), torch.ones(3, 3)))
 
 # Errors
 try:
-    print(aten_dialect(torch.ones(3, 2), torch.ones(3, 2)))
+    print(aten_dialect.module()(torch.ones(3, 2), torch.ones(3, 2)))
 except Exception:
     tb.print_exc()
 
@@ -175,18 +175,18 @@ def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
 # Now let's try running the model with different shapes:
 
 # Works correctly
-print(aten_dialect(torch.ones(3, 3), torch.ones(3, 3)))
-print(aten_dialect(torch.ones(3, 2), torch.ones(3, 2)))
+print(aten_dialect.module()(torch.ones(3, 3), torch.ones(3, 3)))
+print(aten_dialect.module()(torch.ones(3, 2), torch.ones(3, 2)))
 
 # Errors because it violates our constraint that input 0, dim 1 <= 10
 try:
-    print(aten_dialect(torch.ones(3, 15), torch.ones(3, 15)))
+    print(aten_dialect.module()(torch.ones(3, 15), torch.ones(3, 15)))
 except Exception:
     tb.print_exc()
 
 # Errors because it violates our constraint that input 0, dim 1 == input 1, dim 1
 try:
-    print(aten_dialect(torch.ones(3, 3), torch.ones(3, 2)))
+    print(aten_dialect.module()(torch.ones(3, 3), torch.ones(3, 2)))
 except Exception:
     tb.print_exc()
 
@@ -287,23 +287,25 @@ def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
 # there is only one program, it will by default be saved to the name "forward".
 
 
-def encode(x):
-    return torch.nn.functional.linear(x, torch.randn(5, 10))
+class Encode(torch.nn.Module):
+    def forward(self, x):
+        return torch.nn.functional.linear(x, torch.randn(5, 10))
 
 
-def decode(x):
-    return torch.nn.functional.linear(x, torch.randn(10, 5))
+class Decode(torch.nn.Module):
+    def forward(self, x):
+        return torch.nn.functional.linear(x, torch.randn(10, 5))
 
 
 encode_args = (torch.randn(1, 10),)
 aten_encode: ExportedProgram = export(
-    capture_pre_autograd_graph(encode, encode_args),
+    capture_pre_autograd_graph(Encode(), encode_args),
     encode_args,
 )
 
 decode_args = (torch.randn(1, 5),)
 aten_decode: ExportedProgram = export(
-    capture_pre_autograd_graph(decode, decode_args),
+    capture_pre_autograd_graph(Decode(), decode_args),
     decode_args,
 )
 
@@ -486,17 +488,18 @@ def forward(self, x):
 # ``LoweredBackendModule`` for each of those subgraphs.
 
 
-def f(a, x, b):
-    y = torch.mm(a, x)
-    z = y + b
-    a = z - a
-    y = torch.mm(a, x)
-    z = y + b
-    return z
+class Foo(torch.nn.Module):
+    def forward(self, a, x, b):
+        y = torch.mm(a, x)
+        z = y + b
+        a = z - a
+        y = torch.mm(a, x)
+        z = y + b
+        return z
 
 
 example_args = (torch.randn(2, 2), torch.randn(2, 2), torch.randn(2, 2))
-pre_autograd_aten_dialect = capture_pre_autograd_graph(f, example_args)
+pre_autograd_aten_dialect = capture_pre_autograd_graph(Foo(), example_args)
 aten_dialect: ExportedProgram = export(pre_autograd_aten_dialect, example_args)
 edge_program: EdgeProgramManager = to_edge(aten_dialect)
 exported_program = edge_program.exported_program()
@@ -520,17 +523,18 @@ def f(a, x, b):
 # call ``to_backend`` on it:
 
 
-def f(a, x, b):
-    y = torch.mm(a, x)
-    z = y + b
-    a = z - a
-    y = torch.mm(a, x)
-    z = y + b
-    return z
+class Foo(torch.nn.Module):
+    def forward(self, a, x, b):
+        y = torch.mm(a, x)
+        z = y + b
+        a = z - a
+        y = torch.mm(a, x)
+        z = y + b
+        return z
 
 
 example_args = (torch.randn(2, 2), torch.randn(2, 2), torch.randn(2, 2))
-pre_autograd_aten_dialect = capture_pre_autograd_graph(f, example_args)
+pre_autograd_aten_dialect = capture_pre_autograd_graph(Foo(), example_args)
 aten_dialect: ExportedProgram = export(pre_autograd_aten_dialect, example_args)
 edge_program: EdgeProgramManager = to_edge(aten_dialect)
 exported_program = edge_program.exported_program()

examples/apple/coreml/scripts/export.py

@@ -8,19 +8,20 @@
 import pathlib
 import sys
 
-import torch
 import executorch.exir as exir
 
-from executorch.backends.apple.coreml.compiler import CoreMLBackend
+import torch
 
-from executorch.exir.backend.backend_api import to_backend
-from executorch.exir.backend.compile_spec_schema import CompileSpec
-from executorch.sdk.etrecord import generate_etrecord
+from executorch.backends.apple.coreml.compiler import CoreMLBackend
 
 from executorch.backends.apple.coreml.partition.coreml_partitioner import (
     CoreMLPartitioner,
 )
 
+from executorch.exir.backend.backend_api import to_backend
+from executorch.exir.backend.compile_spec_schema import CompileSpec
+from executorch.sdk.etrecord import generate_etrecord
+
 REPO_ROOT = pathlib.Path(__file__).resolve().parent.parent.parent.parent.parent
 EXAMPLES_DIR = REPO_ROOT / "examples"
 sys.path.append(str(EXAMPLES_DIR.absolute()))