[typing] fix register_buffer/parameter (pytorch#40669)

Summary: Pull Request resolved: pytorch#40669

Differential Revision: D22286130

Pulled By: ezyang

fbshipit-source-id: c0cc173279678978726895a0830343d5234e474e

Author: ssnl

torch/nn/modules/module.py

@@ -234,7 +234,7 @@ def _forward_unimplemented(self, *input: Any) -> None:
     """
     forward: Callable[..., Any] = _forward_unimplemented
 
-    def register_buffer(self, name: str, tensor: Tensor, persistent: bool = True) -> None:
+    def register_buffer(self, name: str, tensor: Optional[Tensor], persistent: bool = True) -> None:
         r"""Adds a buffer to the module.
 
         This is typically used to register a buffer that should not to be
@@ -286,7 +286,7 @@ def register_buffer(self, name: str, tensor: Tensor, persistent: bool = True) ->
             else:
                 self._non_persistent_buffers_set.add(name)
 
-    def register_parameter(self, name: str, param: Parameter) -> None:
+    def register_parameter(self, name: str, param: Optional[Parameter]) -> None:
         r"""Adds a parameter to the module.
 
         The parameter can be accessed as an attribute using given name.
@@ -325,7 +325,7 @@ def register_parameter(self, name: str, param: Parameter) -> None:
         else:
             self._parameters[name] = param
 
-    def add_module(self, name: str, module: 'Module') -> None:
+    def add_module(self, name: str, module: Optional['Module']) -> None:
         r"""Adds a child module to the current module.
 
         The module can be accessed as an attribute using the given name.