Fix forward mode ad tutorial by updating stateless import (#2244)

zou3519 · Svetlana Karslioglu · web-flow · commit 9580916266cf · 2023-03-15T13:50:09.000-07:00
Updated the import from the now deleted torch.nn.utils._stateless to torch.func Fixes #2218 Test Plan: - tested locally Co-authored-by: Svetlana Karslioglu <svekars@fb.com>
diff --git a/intermediate_source/forward_ad_usage.py b/intermediate_source/forward_ad_usage.py
@@ -100,13 +100,12 @@ def fn(x, y):
     jvp = fwAD.unpack_dual(out).tangent
 
 ######################################################################
-# Using Modules stateless API (experimental)
+# Using the functional Module API (beta)
 # --------------------------------------------------------------------
 # Another way to use ``nn.Module`` with forward AD is to utilize
-# the stateless API. NB: At the time of writing the stateless API is still
-# experimental and may be subject to change.
+# the functional Module API (also known as the stateless Module API).
 
-from torch.nn.utils._stateless import functional_call
+from torch.func import functional_call
 
 # We need a fresh module because the functional call requires the
 # the model to have parameters registered.
