Updated docstring func signature

pytorch · Apr 25, 2022 · 1340fb0 · 1340fb0
1 parent aa7da3e
commit 1340fb0
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/functorch/_src/make_functional.py b/functorch/_src/make_functional.py
@@ -317,7 +317,7 @@ def forward(self, params, *args, **kwargs):
 
 
 def make_functional(model: nn.Module, disable_autograd_tracking: bool = False):
-    """make_functional(model) -> func, params
+    """make_functional(model, disable_autograd_tracking=False) -> func, params
 
     Given a ``torch.nn.Module``, :func:`make_functional` extracts the state
     (params) and returns a functional version of the model, ``func``. This
@@ -381,7 +381,7 @@ def compute_loss(params, x, t):
 
 
 def make_functional_with_buffers(model: nn.Module, disable_autograd_tracking: bool = False):
-    """make_functional_with_buffers(model) -> func, params, buffers
+    """make_functional_with_buffers(model, disable_autograd_tracking=False) -> func, params, buffers
 
     Given a ``torch.nn.Module``, make_functional_with_buffers extracts the
     state (params and buffers) and returns a functional version of the model