We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
`autoload=True` not working with models in the gpu
If we create the snapshot extension with `autoload=True` the model will not correctly load its state.

`autoload=true` loads the state in the CPU. but it is not executed until the first iteration and it will overwrite the device of the model.  This requires us to call `start_extensions` manually and then do the device move for it to work

```python
# model parameters will be moved to cpu at the beginning of the first iteration due autoload being executed there
...
manager.extend(extensions.snapshot(autoload=True)
# move the model, but this will be overwriten later
model.cuda()

for batch in train_loader:
    with manager.run_iteration():   # Snapshot load happens the first time this is executed
        model(batch.cuda())   #  Error! weights are in the cpu again
```

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `pytorch_pfn_extras/nn/modules/lazy.py`
Content:
```
1 import inspect
2 from typing import Tuple
3 import warnings
4 
5 import torch
6 
7 
8 class LazyInitializationMixin:
9 
10     """A mixin for modules that lazily initialize buffers and parameters.
11 
12     Unlike regular modules, subclasses of this module can initialize
13     buffers and parameters outside of the constructor (``__init__``).
14     This allows you to, for example, initialize parameters in ``forward``
15     method to determine the shape of the weight based on the initial input.
16 
17     Be sure to run "dummy" forward once to initialize all parameters that
18     should be trained, before passing ``module.parameters()`` to an optimizer;
19     otherwise weights initialized after ``module.parameters()`` (e.g., in
20     ``forward`` function) will never be trained.
21 
22     Note that lazy modules cannot validate if the shape is correct during
23     deserialization.  Also note that the initial weights may become different
24     from the original (non-lazy) module even if the random seed is manually
25     configured, as the order of initialization is different from the original
26     one; especially, ``module.cuda()`` may cause the initialization to run on
27     a GPU.
28 
29     The default value of lazy buffers and parameters are ``torch.Tensor([])``
30     and ``UninitializedParameter()``, respectively.
31     """
32 
33     # Subclasses must override these fields and list names of all buffers /
34     # parameters that will be initialized lazily.
35     lazy_buffer_names: Tuple[str, ...] = ()
36     lazy_parameter_names: Tuple[str, ...] = ()
37 
38     def __init__(self, *args, **kwargs):
39         self._lazy_ready = False
40 
41         super().__init__(*args, **kwargs)
42 
43         for name in self.lazy_buffer_names:
44             self.register_buffer(name, torch.Tensor([]))
45         for name in self.lazy_parameter_names:
46             self.register_parameter(name, UninitializedParameter())
47         self._register_load_state_dict_pre_hook(self._lazy_load_hook)
48         self._lazy_ready = True
49 
50     @property
51     def lazy_parmeters_determined(self):
52         """Returns if all lazy parameters are determined.
53 
54         Subclasses can perform parameters initialization after all lazy
55         parameters are determined.  Note that this may be called during
56         ``__init__``.
57         """
58         return self._lazy_ready and all([
59             not isinstance(getattr(self, x), UninitializedParameter)
60             for x in self.lazy_parameter_names])
61 
62     def state_dict(self, *args, **kwargs):
63         """Returns a dictionary containing a whole state of the module.
64 
65         This function overrides the default behavior to exclude uninitialized
66         parameter from serialization.  This is needed because we need to
67         discriminate lazy parameters (``UninitializedParameter()`) and
68         initialized empty parameters (``torch.nn.Parameter(torch.Tensor())``)
69         during deserialization.
70 
71         See comments of ``_lazy_load_hook`` for details.
72         """
73         destination = super().state_dict(*args, **kwargs)
74         for name in self.lazy_parameter_names:
75             if isinstance(getattr(self, name), UninitializedParameter):
76                 del destination[name]
77         return destination
78 
79     def _lazy_load_hook(
80             self, state_dict, prefix, local_metadata, strict,
81             missing_keys, unexpected_keys, error_msgs):
82         """load_state_dict pre-hook function for lazy buffers and parameters.
83 
84         The purpose of this hook is to adjust the current state and/or
85         ``state_dict`` being loaded so that a module instance serialized in
86         both un/initialized state can be deserialized onto both un/initialized
87         module instance.
88 
89         See comment in ``torch.nn.Module._register_load_state_dict_pre_hook``
90         for the details of the hook specification.
91         """
92         for name in self.lazy_buffer_names:
93             # Avoid shape mismatch error when loading an initialized buffer
94             # onto an uninitialized module instance.
95             self.register_buffer(name, state_dict[prefix + name])
96 
97         for name in self.lazy_parameter_names:
98             # The parameter may not exist in the loaded ``state_dict`` if the
99             # original module was serialized before initializing lazy
100             # parameters (see comments of ``state_dict``).
101             key = prefix + name
102             if key in state_dict:
103                 # The model was serialized after initialization.
104                 self.register_parameter(
105                     name, torch.nn.Parameter(state_dict[key]))
106             else:
107                 # The model was serialized before initialization.
108                 param = UninitializedParameter()
109                 self.register_parameter(name, param)
110                 state_dict[key] = param
111 
112 
113 class UninitializedParameter(torch.nn.Parameter):
114 
115     def __repr__(self):
116         return 'Uninitialized lazy parameter'
117 
118     def share_memory_(self):
119         raise RuntimeError(
120             'Can\'t share memory on an unitialized parameter. '
121             'Run forward to initialize the network before calling '
122             '`module.share_memory()`.')
123 
124     @property
125     def is_leaf(self):
126         # Hacky workaround to detect use of uninitialized lazy parameters.
127         # This overrides ``is_leaf`` attribute which should always be ``True``
128         # for parameters; optimizers check for this attribute and raise an
129         # error if non-leaf tensors are detected.
130         frame = inspect.currentframe()
131         if frame.f_back.f_globals['__package__'].startswith('torch.optim'):
132             warnings.warn('''
133     Use of uninitialized lazy parameter in Optimizer has been detected.
134     Maybe you forgot to run forward before passing `module.parameters()` to the optimizer?''')  # NOQA
135         return True
136
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/pytorch_pfn_extras/nn/modules/lazy.py b/pytorch_pfn_extras/nn/modules/lazy.py
--- a/pytorch_pfn_extras/nn/modules/lazy.py
+++ b/pytorch_pfn_extras/nn/modules/lazy.py
@@ -81,32 +81,45 @@
             missing_keys, unexpected_keys, error_msgs):
         """load_state_dict pre-hook function for lazy buffers and parameters.
 
-        The purpose of this hook is to adjust the current state and/or
-        ``state_dict`` being loaded so that a module instance serialized in
-        both un/initialized state can be deserialized onto both un/initialized
-        module instance.
+        The purpose of this hook is to check the current state and/or
+        ``state_dict`` being loaded and ensure that both are states
+        are properly initialized.
 
         See comment in ``torch.nn.Module._register_load_state_dict_pre_hook``
         for the details of the hook specification.
         """
         for name in self.lazy_buffer_names:
-            # Avoid shape mismatch error when loading an initialized buffer
-            # onto an uninitialized module instance.
-            self.register_buffer(name, state_dict[prefix + name])
+            key = prefix + name
+            module_initialized = getattr(self, name).shape != (0,)
+            state_initialized = state_dict[key].shape != (0,)
+            if module_initialized and not state_initialized:
+                raise RuntimeError(
+                    'Can\'t load non-initialized buffers in already '
+                    'initialized modules')
+            elif not module_initialized and state_initialized:
+                # Here we need to avoid a tensor size mismatch
+                # this is a regular tensor without a materialize
+                # method, so we can just resize for the load logic to copy
+                # the contents later to the correct device the module
+                # was moved to
+                getattr(self, name).resize_(state_dict[key].size())
 
         for name in self.lazy_parameter_names:
-            # The parameter may not exist in the loaded ``state_dict`` if the
+            # The parameter does not exist in the loaded ``state_dict`` if the
             # original module was serialized before initializing lazy
             # parameters (see comments of ``state_dict``).
             key = prefix + name
-            if key in state_dict:
-                # The model was serialized after initialization.
-                self.register_parameter(
-                    name, torch.nn.Parameter(state_dict[key]))
-            else:
-                # The model was serialized before initialization.
+            module_initialized = not isinstance(
+                getattr(self, name), UninitializedParameter)
+            state_initialized = key in state_dict
+            if module_initialized and not state_initialized:
+                raise RuntimeError(
+                    'Can\'t load uninitialized parameters in already '
+                    'initialized modules')
+            elif not module_initialized and state_initialized:
+                getattr(self, name).materialize(state_dict[key].shape)
+            elif key not in state_dict and not module_initialized:
                 param = UninitializedParameter()
-                self.register_parameter(name, param)
                 state_dict[key] = param
 
 
@@ -133,3 +146,25 @@
     Use of uninitialized lazy parameter in Optimizer has been detected.
     Maybe you forgot to run forward before passing `module.parameters()` to the optimizer?''')  # NOQA
         return True
+
+    def materialize(self, shape, device=None, dtype=None):
+        r"""Create a Parameter with the same properties of the uninitialized
+        one. Given a shape, it materializes a parameter in the same device
+        and with the same `dtype` as the current one or the specified ones in
+        the arguments.
+
+        Args:
+            shape : (tuple): the shape for the materialized tensor.
+            device (:class:`torch.device`): the desired device of the
+                parameters
+                and buffers in this module. Optional.
+            dtype (:class:`torch.dtype`): the desired floating point type of
+                the floating point parameters and buffers in this module.
+                Optional.
+        """
+        if device is None:
+            device = self.data.device
+        if dtype is None:
+            dtype = self.data.dtype
+        self.data = torch.empty(shape, device=device, dtype=dtype)
+        self.__class__ = torch.nn.Parameter

{"golden_diff": "diff --git a/pytorch_pfn_extras/nn/modules/lazy.py b/pytorch_pfn_extras/nn/modules/lazy.py\n--- a/pytorch_pfn_extras/nn/modules/lazy.py\n+++ b/pytorch_pfn_extras/nn/modules/lazy.py\n@@ -81,32 +81,45 @@\n             missing_keys, unexpected_keys, error_msgs):\n         \"\"\"load_state_dict pre-hook function for lazy buffers and parameters.\n \n-        The purpose of this hook is to adjust the current state and/or\n-        ``state_dict`` being loaded so that a module instance serialized in\n-        both un/initialized state can be deserialized onto both un/initialized\n-        module instance.\n+        The purpose of this hook is to check the current state and/or\n+        ``state_dict`` being loaded and ensure that both are states\n+        are properly initialized.\n \n         See comment in ``torch.nn.Module._register_load_state_dict_pre_hook``\n         for the details of the hook specification.\n         \"\"\"\n         for name in self.lazy_buffer_names:\n-            # Avoid shape mismatch error when loading an initialized buffer\n-            # onto an uninitialized module instance.\n-            self.register_buffer(name, state_dict[prefix + name])\n+            key = prefix + name\n+            module_initialized = getattr(self, name).shape != (0,)\n+            state_initialized = state_dict[key].shape != (0,)\n+            if module_initialized and not state_initialized:\n+                raise RuntimeError(\n+                    'Can\\'t load non-initialized buffers in already '\n+                    'initialized modules')\n+            elif not module_initialized and state_initialized:\n+                # Here we need to avoid a tensor size mismatch\n+                # this is a regular tensor without a materialize\n+                # method, so we can just resize for the load logic to copy\n+                # the contents later to the correct device the module\n+                # was moved to\n+                getattr(self, name).resize_(state_dict[key].size())\n \n         for name in self.lazy_parameter_names:\n-            # The parameter may not exist in the loaded ``state_dict`` if the\n+            # The parameter does not exist in the loaded ``state_dict`` if the\n             # original module was serialized before initializing lazy\n             # parameters (see comments of ``state_dict``).\n             key = prefix + name\n-            if key in state_dict:\n-                # The model was serialized after initialization.\n-                self.register_parameter(\n-                    name, torch.nn.Parameter(state_dict[key]))\n-            else:\n-                # The model was serialized before initialization.\n+            module_initialized = not isinstance(\n+                getattr(self, name), UninitializedParameter)\n+            state_initialized = key in state_dict\n+            if module_initialized and not state_initialized:\n+                raise RuntimeError(\n+                    'Can\\'t load uninitialized parameters in already '\n+                    'initialized modules')\n+            elif not module_initialized and state_initialized:\n+                getattr(self, name).materialize(state_dict[key].shape)\n+            elif key not in state_dict and not module_initialized:\n                 param = UninitializedParameter()\n-                self.register_parameter(name, param)\n                 state_dict[key] = param\n \n \n@@ -133,3 +146,25 @@\n     Use of uninitialized lazy parameter in Optimizer has been detected.\n     Maybe you forgot to run forward before passing `module.parameters()` to the optimizer?''')  # NOQA\n         return True\n+\n+    def materialize(self, shape, device=None, dtype=None):\n+        r\"\"\"Create a Parameter with the same properties of the uninitialized\n+        one. Given a shape, it materializes a parameter in the same device\n+        and with the same `dtype` as the current one or the specified ones in\n+        the arguments.\n+\n+        Args:\n+            shape : (tuple): the shape for the materialized tensor.\n+            device (:class:`torch.device`): the desired device of the\n+                parameters\n+                and buffers in this module. Optional.\n+            dtype (:class:`torch.dtype`): the desired floating point type of\n+                the floating point parameters and buffers in this module.\n+                Optional.\n+        \"\"\"\n+        if device is None:\n+            device = self.data.device\n+        if dtype is None:\n+            dtype = self.data.dtype\n+        self.data = torch.empty(shape, device=device, dtype=dtype)\n+        self.__class__ = torch.nn.Parameter\n", "issue": "`autoload=True` not working with models in the gpu\nIf we create the snapshot extension with `autoload=True` the model will not correctly load its state.\r\n\r\n`autoload=true` loads the state in the CPU. but it is not executed until the first iteration and it will overwrite the device of the model.  This requires us to call `start_extensions` manually and then do the device move for it to work\r\n\r\n```python\r\n# model parameters will be moved to cpu at the beginning of the first iteration due autoload being executed there\r\n...\r\nmanager.extend(extensions.snapshot(autoload=True)\r\n# move the model, but this will be overwriten later\r\nmodel.cuda()\r\n\r\nfor batch in train_loader:\r\n    with manager.run_iteration():   # Snapshot load happens the first time this is executed\r\n        model(batch.cuda())   #  Error! weights are in the cpu again\r\n```\n", "before_files": [{"content": "import inspect\nfrom typing import Tuple\nimport warnings\n\nimport torch\n\n\nclass LazyInitializationMixin:\n\n    \"\"\"A mixin for modules that lazily initialize buffers and parameters.\n\n    Unlike regular modules, subclasses of this module can initialize\n    buffers and parameters outside of the constructor (``__init__``).\n    This allows you to, for example, initialize parameters in ``forward``\n    method to determine the shape of the weight based on the initial input.\n\n    Be sure to run \"dummy\" forward once to initialize all parameters that\n    should be trained, before passing ``module.parameters()`` to an optimizer;\n    otherwise weights initialized after ``module.parameters()`` (e.g., in\n    ``forward`` function) will never be trained.\n\n    Note that lazy modules cannot validate if the shape is correct during\n    deserialization.  Also note that the initial weights may become different\n    from the original (non-lazy) module even if the random seed is manually\n    configured, as the order of initialization is different from the original\n    one; especially, ``module.cuda()`` may cause the initialization to run on\n    a GPU.\n\n    The default value of lazy buffers and parameters are ``torch.Tensor([])``\n    and ``UninitializedParameter()``, respectively.\n    \"\"\"\n\n    # Subclasses must override these fields and list names of all buffers /\n    # parameters that will be initialized lazily.\n    lazy_buffer_names: Tuple[str, ...] = ()\n    lazy_parameter_names: Tuple[str, ...] = ()\n\n    def __init__(self, *args, **kwargs):\n        self._lazy_ready = False\n\n        super().__init__(*args, **kwargs)\n\n        for name in self.lazy_buffer_names:\n            self.register_buffer(name, torch.Tensor([]))\n        for name in self.lazy_parameter_names:\n            self.register_parameter(name, UninitializedParameter())\n        self._register_load_state_dict_pre_hook(self._lazy_load_hook)\n        self._lazy_ready = True\n\n    @property\n    def lazy_parmeters_determined(self):\n        \"\"\"Returns if all lazy parameters are determined.\n\n        Subclasses can perform parameters initialization after all lazy\n        parameters are determined.  Note that this may be called during\n        ``__init__``.\n        \"\"\"\n        return self._lazy_ready and all([\n            not isinstance(getattr(self, x), UninitializedParameter)\n            for x in self.lazy_parameter_names])\n\n    def state_dict(self, *args, **kwargs):\n        \"\"\"Returns a dictionary containing a whole state of the module.\n\n        This function overrides the default behavior to exclude uninitialized\n        parameter from serialization.  This is needed because we need to\n        discriminate lazy parameters (``UninitializedParameter()`) and\n        initialized empty parameters (``torch.nn.Parameter(torch.Tensor())``)\n        during deserialization.\n\n        See comments of ``_lazy_load_hook`` for details.\n        \"\"\"\n        destination = super().state_dict(*args, **kwargs)\n        for name in self.lazy_parameter_names:\n            if isinstance(getattr(self, name), UninitializedParameter):\n                del destination[name]\n        return destination\n\n    def _lazy_load_hook(\n            self, state_dict, prefix, local_metadata, strict,\n            missing_keys, unexpected_keys, error_msgs):\n        \"\"\"load_state_dict pre-hook function for lazy buffers and parameters.\n\n        The purpose of this hook is to adjust the current state and/or\n        ``state_dict`` being loaded so that a module instance serialized in\n        both un/initialized state can be deserialized onto both un/initialized\n        module instance.\n\n        See comment in ``torch.nn.Module._register_load_state_dict_pre_hook``\n        for the details of the hook specification.\n        \"\"\"\n        for name in self.lazy_buffer_names:\n            # Avoid shape mismatch error when loading an initialized buffer\n            # onto an uninitialized module instance.\n            self.register_buffer(name, state_dict[prefix + name])\n\n        for name in self.lazy_parameter_names:\n            # The parameter may not exist in the loaded ``state_dict`` if the\n            # original module was serialized before initializing lazy\n            # parameters (see comments of ``state_dict``).\n            key = prefix + name\n            if key in state_dict:\n                # The model was serialized after initialization.\n                self.register_parameter(\n                    name, torch.nn.Parameter(state_dict[key]))\n            else:\n                # The model was serialized before initialization.\n                param = UninitializedParameter()\n                self.register_parameter(name, param)\n                state_dict[key] = param\n\n\nclass UninitializedParameter(torch.nn.Parameter):\n\n    def __repr__(self):\n        return 'Uninitialized lazy parameter'\n\n    def share_memory_(self):\n        raise RuntimeError(\n            'Can\\'t share memory on an unitialized parameter. '\n            'Run forward to initialize the network before calling '\n            '`module.share_memory()`.')\n\n    @property\n    def is_leaf(self):\n        # Hacky workaround to detect use of uninitialized lazy parameters.\n        # This overrides ``is_leaf`` attribute which should always be ``True``\n        # for parameters; optimizers check for this attribute and raise an\n        # error if non-leaf tensors are detected.\n        frame = inspect.currentframe()\n        if frame.f_back.f_globals['__package__'].startswith('torch.optim'):\n            warnings.warn('''\n    Use of uninitialized lazy parameter in Optimizer has been detected.\n    Maybe you forgot to run forward before passing `module.parameters()` to the optimizer?''')  # NOQA\n        return True\n", "path": "pytorch_pfn_extras/nn/modules/lazy.py"}], "after_files": [{"content": "import inspect\nfrom typing import Tuple\nimport warnings\n\nimport torch\n\n\nclass LazyInitializationMixin:\n\n    \"\"\"A mixin for modules that lazily initialize buffers and parameters.\n\n    Unlike regular modules, subclasses of this module can initialize\n    buffers and parameters outside of the constructor (``__init__``).\n    This allows you to, for example, initialize parameters in ``forward``\n    method to determine the shape of the weight based on the initial input.\n\n    Be sure to run \"dummy\" forward once to initialize all parameters that\n    should be trained, before passing ``module.parameters()`` to an optimizer;\n    otherwise weights initialized after ``module.parameters()`` (e.g., in\n    ``forward`` function) will never be trained.\n\n    Note that lazy modules cannot validate if the shape is correct during\n    deserialization.  Also note that the initial weights may become different\n    from the original (non-lazy) module even if the random seed is manually\n    configured, as the order of initialization is different from the original\n    one; especially, ``module.cuda()`` may cause the initialization to run on\n    a GPU.\n\n    The default value of lazy buffers and parameters are ``torch.Tensor([])``\n    and ``UninitializedParameter()``, respectively.\n    \"\"\"\n\n    # Subclasses must override these fields and list names of all buffers /\n    # parameters that will be initialized lazily.\n    lazy_buffer_names: Tuple[str, ...] = ()\n    lazy_parameter_names: Tuple[str, ...] = ()\n\n    def __init__(self, *args, **kwargs):\n        self._lazy_ready = False\n\n        super().__init__(*args, **kwargs)\n\n        for name in self.lazy_buffer_names:\n            self.register_buffer(name, torch.Tensor([]))\n        for name in self.lazy_parameter_names:\n            self.register_parameter(name, UninitializedParameter())\n        self._register_load_state_dict_pre_hook(self._lazy_load_hook)\n        self._lazy_ready = True\n\n    @property\n    def lazy_parmeters_determined(self):\n        \"\"\"Returns if all lazy parameters are determined.\n\n        Subclasses can perform parameters initialization after all lazy\n        parameters are determined.  Note that this may be called during\n        ``__init__``.\n        \"\"\"\n        return self._lazy_ready and all([\n            not isinstance(getattr(self, x), UninitializedParameter)\n            for x in self.lazy_parameter_names])\n\n    def state_dict(self, *args, **kwargs):\n        \"\"\"Returns a dictionary containing a whole state of the module.\n\n        This function overrides the default behavior to exclude uninitialized\n        parameter from serialization.  This is needed because we need to\n        discriminate lazy parameters (``UninitializedParameter()`) and\n        initialized empty parameters (``torch.nn.Parameter(torch.Tensor())``)\n        during deserialization.\n\n        See comments of ``_lazy_load_hook`` for details.\n        \"\"\"\n        destination = super().state_dict(*args, **kwargs)\n        for name in self.lazy_parameter_names:\n            if isinstance(getattr(self, name), UninitializedParameter):\n                del destination[name]\n        return destination\n\n    def _lazy_load_hook(\n            self, state_dict, prefix, local_metadata, strict,\n            missing_keys, unexpected_keys, error_msgs):\n        \"\"\"load_state_dict pre-hook function for lazy buffers and parameters.\n\n        The purpose of this hook is to check the current state and/or\n        ``state_dict`` being loaded and ensure that both are states\n        are properly initialized.\n\n        See comment in ``torch.nn.Module._register_load_state_dict_pre_hook``\n        for the details of the hook specification.\n        \"\"\"\n        for name in self.lazy_buffer_names:\n            key = prefix + name\n            module_initialized = getattr(self, name).shape != (0,)\n            state_initialized = state_dict[key].shape != (0,)\n            if module_initialized and not state_initialized:\n                raise RuntimeError(\n                    'Can\\'t load non-initialized buffers in already '\n                    'initialized modules')\n            elif not module_initialized and state_initialized:\n                # Here we need to avoid a tensor size mismatch\n                # this is a regular tensor without a materialize\n                # method, so we can just resize for the load logic to copy\n                # the contents later to the correct device the module\n                # was moved to\n                getattr(self, name).resize_(state_dict[key].size())\n\n        for name in self.lazy_parameter_names:\n            # The parameter does not exist in the loaded ``state_dict`` if the\n            # original module was serialized before initializing lazy\n            # parameters (see comments of ``state_dict``).\n            key = prefix + name\n            module_initialized = not isinstance(\n                getattr(self, name), UninitializedParameter)\n            state_initialized = key in state_dict\n            if module_initialized and not state_initialized:\n                raise RuntimeError(\n                    'Can\\'t load uninitialized parameters in already '\n                    'initialized modules')\n            elif not module_initialized and state_initialized:\n                getattr(self, name).materialize(state_dict[key].shape)\n            elif key not in state_dict and not module_initialized:\n                param = UninitializedParameter()\n                state_dict[key] = param\n\n\nclass UninitializedParameter(torch.nn.Parameter):\n\n    def __repr__(self):\n        return 'Uninitialized lazy parameter'\n\n    def share_memory_(self):\n        raise RuntimeError(\n            'Can\\'t share memory on an unitialized parameter. '\n            'Run forward to initialize the network before calling '\n            '`module.share_memory()`.')\n\n    @property\n    def is_leaf(self):\n        # Hacky workaround to detect use of uninitialized lazy parameters.\n        # This overrides ``is_leaf`` attribute which should always be ``True``\n        # for parameters; optimizers check for this attribute and raise an\n        # error if non-leaf tensors are detected.\n        frame = inspect.currentframe()\n        if frame.f_back.f_globals['__package__'].startswith('torch.optim'):\n            warnings.warn('''\n    Use of uninitialized lazy parameter in Optimizer has been detected.\n    Maybe you forgot to run forward before passing `module.parameters()` to the optimizer?''')  # NOQA\n        return True\n\n    def materialize(self, shape, device=None, dtype=None):\n        r\"\"\"Create a Parameter with the same properties of the uninitialized\n        one. Given a shape, it materializes a parameter in the same device\n        and with the same `dtype` as the current one or the specified ones in\n        the arguments.\n\n        Args:\n            shape : (tuple): the shape for the materialized tensor.\n            device (:class:`torch.device`): the desired device of the\n                parameters\n                and buffers in this module. Optional.\n            dtype (:class:`torch.dtype`): the desired floating point type of\n                the floating point parameters and buffers in this module.\n                Optional.\n        \"\"\"\n        if device is None:\n            device = self.data.device\n        if dtype is None:\n            dtype = self.data.dtype\n        self.data = torch.empty(shape, device=device, dtype=dtype)\n        self.__class__ = torch.nn.Parameter\n", "path": "pytorch_pfn_extras/nn/modules/lazy.py"}]}