Update meta_init_context.py

meta_init_context.py

@@ -1,4 +1,5 @@
 from contextlib import contextmanager
+from typing import Any, Callable, Optional
 import torch
 import torch.nn as nn
 
@@ -57,25 +58,29 @@ def init_on_device(device: torch.device, include_buffers: bool=False):
     if include_buffers:
         old_register_buffer = nn.Module.register_buffer
 
-    def register_empty_parameter(module, name, param):
-        old_register_parameter(module, name, param)
+    def register_empty_parameter(self: torch.nn.Module, name: str, param: Optional[torch.nn.Parameter]):
+        old_register_parameter(self, name, param)
         if param is not None:
-            param_cls = type(module._parameters[name])
-            kwargs = module._parameters[name].__dict__
-            module._parameters[name] = param_cls(module._parameters[name].to(device), **kwargs)
-
-    def register_empty_buffer(module, name, buffer):
-        old_register_buffer(module, name, buffer)
-        if buffer is not None:
-            module._buffers[name] = module._buffers[name].to(device)
+            parameter = self._parameters[name]
+            assert parameter is not None
+            param_cls = type(parameter)
+            kwargs = parameter.__dict__
+            self._parameters[name] = param_cls(parameter.to(device), **kwargs)
+
+    def register_empty_buffer(self: torch.nn.Module, name: str, tensor: Optional[torch.Tensor], persistent: bool=True):
+        old_register_buffer(self, name, tensor, persistent=persistent)
+        if tensor is not None:
+            named_buffer = self._buffers[name]
+            assert named_buffer is not None
+            self._buffers[name] = named_buffer.to(device)
     if include_buffers:
         tensor_constructors_to_patch = {torch_function_name: getattr(torch, torch_function_name) for torch_function_name in ['empty', 'zeros', 'ones', 'full']}
     else:
         tensor_constructors_to_patch = {}
 
-    def patch_tensor_constructor(fn):
+    def patch_tensor_constructor(fn: Callable):
 
-        def wrapper(*args, **kwargs):
+        def wrapper(*args: Any, **kwargs: Any):
             kwargs['device'] = device
             return fn(*args, **kwargs)
         return wrapper