class_docstr: Implementation of PureFunction for torch.nn.Module.
  
  Examples
  --------
  >>> import torch
  >>> from deepchem.utils.differentiation_utils import get_pure_function
  >>> class A(torch.nn.Module):
  ...     def __init__(self, a):
  ...         super().__init__()
  ...         self.b = torch.nn.Parameter(torch.tensor(a*a))
  ...     def forward(self, x):
  ...         return self.b * x
  >>> B = A(4.)
  >>> m = get_pure_function(B.forward)
  >>> m.set_objparams([3.])
  >>> m(2)
  6.0
class_docstr_tok: ["Implementation","of","PureFunction","for","torch.nn.Module",".","Examples","--","--","--","--",">",">",">","import","torch",">",">",">","from","deepchem.utils.differentiation_utils","import","get_pure_function",">",">",">","class","A","(","torch.nn.Module",")",":","...","def","__init__","(","self",",","a",")",":","...","super","(",")",".__init__","(",")","...","self.b","=","torch.nn.Parameter","(","torch.tensor","(","a","*","a",")",")","...","def","forward","(","self",",","x",")",":","...","return","self.b","*","x",">",">",">","B","=","A","(","4",".",")",">",">",">","m","=","get_pure_function","(","B.forward",")",">",">",">","m.set_objparams","(","[","3",".","]",")",">",">",">","m","(","2",")","6.0"]
docstr: Initialize the TorchNNPureFunction.
  
  Parameters
  ----------
  obj: torch.nn.Module
      Object to be wrapped
  method: Callable
      Method to be wrapped
docstr_tok: ["Initialize","the","TorchNNPureFunction",".","Parameters","--","--","--","--","--","obj",":","torch.nn.Module","Object","to","be","wrapped","method",":","Callable","Method","to","be","wrapped"]
code: def __init__(self, obj: torch.nn.Module, method: Callable):
      """Initialize the TorchNNPureFunction.
  
      Parameters
      ----------
      obj: torch.nn.Module
          Object to be wrapped
      method: Callable
          Method to be wrapped
  
      """
      self.obj = obj
      self.method = method
      super().__init__(method)
  
imports: import torch
  import inspect
  from typing import Callable, List, Tuple, Union
  from deepchem.utils.attribute_utils import set_attr, del_attr
  from deepchem.utils.differentiation_utils import EditableModule
  from deepchem.utils.misc_utils import Uniquifier
  from contextlib import contextmanager
  from abc import abstractmethod
label_desc: Use image node_id 1 to create a new TorchNNPureFunction object from inherited base classes: PureFunction with example: obj = TorchNNPureFunction(obj, method)