Add test_make_fx_model_train example (#980)

Summary: Pull Request resolved: #980

Differential Revision: D38078694

Pulled By: mostafaelhoushi

fbshipit-source-id: 4e8be4c653ac2c871d2e47064644b9a965e05187

Author: mostafaelhoushi

test/test_pythonkey.py

@@ -11,6 +11,7 @@
 import unittest
 import warnings
 import itertools
+import torch.nn.utils._stateless as stateless
 from functools import partial
 from torch.testing._internal.common_device_type import instantiate_device_type_tests
 from functorch import (
@@ -74,6 +75,26 @@ def f(x):
         new_inp = torch.randn(3)
         self.assertEqual(fx_f(new_inp), f(new_inp))
 
+    def test_make_fx_model_train(self, device):
+        class Foo(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.linear = nn.Linear(5, 5)
+
+            def forward(self, x):
+                return self.linear(x).relu()
+        mod = Foo()
+        optim = torch.optim.SGD(mod.parameters(), lr=1e-4)
+
+        def f(x, params):
+            out = stateless.functional_call(mod, params, x).sum()
+            out.backward()
+            return list(params.values())
+        inp = torch.randn(3, 5, requires_grad=True)
+        params = dict(mod.named_parameters())
+        fx_f = make_fx(f)(inp, params)
+        self.assertEqual(fx_f(inp, params), f(inp, params))
+
     def test_scalar_device(self, device):
         def f(a, b):
             return a + b