What the object holds

Construct nn.Module() with no subclass, no layers, nothing, and its __dict__ already has seventeen keys. Print them and the shape of the class shows itself: _parameters, _buffers and _modules are the three registries, _non_persistent_buffers_set marks which buffers stay out of a checkpoint, training is the flag chapter 3 explains, and the remaining eleven are hook tables plus one flag about which kind of backward hook is in use.

The way those slots get created is worth reading, because the comment above them explains the whole design in four lines. Every one is written with super().__setattr__, deliberately bypassing the module's own attribute assignment, since that assignment is the thing being set up here and calling it before _parameters exists would fail.

That failure is real and you can trigger it. Assign an nn.Parameter to self before calling super().__init__() and you get AttributeError: cannot assign parameters before Module.__init__() call, raised because self.__dict__.get('_parameters') came back None.

        """
        Calls super().__setattr__('a', a) instead of the typical self.a = a
        to avoid Module.__setattr__ overhead. Module's __setattr__ has special
        handling for parameters, submodules, and buffers but simply calls into
        super().__setattr__ for all other attributes.
        """
        super().__setattr__('training', True)
        super().__setattr__('_parameters', OrderedDict())
        super().__setattr__('_buffers', OrderedDict())
        super().__setattr__('_non_persistent_buffers_set', set())
        super().__setattr__('_backward_pre_hooks', OrderedDict())
        super().__setattr__('_backward_hooks', OrderedDict())
        super().__setattr__('_is_full_backward_hook', None)
        super().__setattr__('_forward_hooks', OrderedDict())
        super().__setattr__('_forward_hooks_with_kwargs', OrderedDict())
        super().__setattr__('_forward_hooks_always_called', OrderedDict())
        super().__setattr__('_forward_pre_hooks', OrderedDict())
        super().__setattr__('_forward_pre_hooks_with_kwargs', OrderedDict())
        super().__setattr__('_state_dict_hooks', OrderedDict())
        super().__setattr__('_state_dict_pre_hooks', OrderedDict())
        super().__setattr__('_load_state_dict_pre_hooks', OrderedDict())
        super().__setattr__('_load_state_dict_post_hooks', OrderedDict())
        super().__setattr__('_modules', OrderedDict())

self.up = nn.Linear(8, 32) does not store a reference. It calls Module.__setattr__, which runs a chain of type tests and routes the value into one of four places, and the order of those tests is the rule you actually need.

A Parameter goes first: it gets removed from the plain __dict__ and from the other two registries, then registered. Next comes a guard: if the name is already a parameter and the new value is not a Parameter, you get a TypeError rather than a silent downgrade to a plain attribute. Modules and buffers follow the same two-step shape, each with its own guard.

The last branch is where most bugs live, because it does nothing special at all. Anything that is not a Parameter, not a Module, and not landing on an existing buffer name falls through to super().__setattr__(name, value), which is ordinary Python attribute storage. A tensor put there is not part of the model in any sense the rest of the library recognises.

        params = self.__dict__.get('_parameters')
        if isinstance(value, Parameter):
            if params is None:
                raise AttributeError(
                    "cannot assign parameters before Module.__init__() call")
            remove_from(self.__dict__, self._buffers, self._modules, self._non_persistent_buffers_set)
            self.register_parameter(name, value)
        elif params is not None and name in params:
            if value is not None:
                raise TypeError(f"cannot assign '{torch.typename(value)}' as parameter '{name}' "
                                "(torch.nn.Parameter or None expected)"
                                )
            self.register_parameter(name, value)

Five assignments, five different fates. The run below registers a parameter, a persistent buffer, a non-persistent buffer, a plain tensor and a submodule, then asks the three questions that matter: does it train, does it get saved, does it move.

The plain tensor answers no to all three. It never appears in named_parameters() or state_dict(), and after m.double() it is still float32 while the parameter beside it is float64, because the walker that performs a conversion visits _parameters and _buffers and nothing else. A constant you keep on self and use inside forward will therefore be on the wrong device the first time you run on an accelerator.

The fix is not to remember harder. Register it, or build it inside forward from something that is registered, so the answer stops depending on anyone noticing.

Registration is the only thing that makes a tensor part of the model. Assignment on its own is just Python.

import torch
from torch import nn


class Holder(nn.Module):
    def __init__(self):
        super().__init__()
        self.w = nn.Parameter(torch.zeros(2))
        self.register_buffer("steps", torch.zeros(1))
        self.register_buffer("cache", torch.zeros(2), persistent=False)
        self.raw = torch.ones(2)                 # a plain attribute
        self.inner = nn.Linear(2, 2)


m = Holder()
print(list(dict(m.named_parameters())))   # ['w', 'inner.weight', 'inner.bias']
print(list(dict(m.named_buffers())))      # ['steps', 'cache']
print(list(m.state_dict()))               # ['w', 'steps', 'inner.weight', 'inner.bias']
m.double()
print(m.w.dtype, m.raw.dtype)             # torch.float64 torch.float32

persistent=False is the one registration that says yes to training-time behaviour and no to serialization. The buffer moves with the module, shows up in buffers(), and is simply absent from state_dict(), which is what you want for a mask or a cached table you can rebuild from the shapes.

Flipping that flag on an existing model has a consequence people meet as a mystery. A checkpoint written before the flip still carries the key, and loading it into the new class reports the key as unexpected rather than ignoring it, because the load walks the module's own persistent names and treats everything else in the dictionary as surplus.

Under strict=True that surplus is a RuntimeError naming the key. It is the same machinery that catches a genuine typo, which is why it fires on a change you made on purpose.

import torch
from torch import nn


class Old(nn.Module):
    def __init__(self):
        super().__init__()
        self.register_buffer("mask", torch.ones(3))


class New(nn.Module):
    def __init__(self):
        super().__init__()
        self.register_buffer("mask", torch.ones(3), persistent=False)


ck = Old().state_dict()
print(list(New().buffers()))                       # [tensor([1., 1., 1.])]
print(New().load_state_dict(ck, strict=False))
# _IncompatibleKeys(missing_keys=[], unexpected_keys=['mask'])
New().load_state_dict(ck)
# RuntimeError: Error(s) in loading state_dict for New:
#     Unexpected key(s) in state_dict: "mask".

Back to the eleven hook dictionaries. A hook is a callable you register on one module instance; the registration returns a handle whose remove() deletes that one entry. Forward pre-hooks run before forward, forward hooks run after it and may replace the output by returning a value, and backward hooks run when the gradient reaches the module.

Two of those tables share a name for a reason. register_backward_hook and register_full_backward_hook both write into _backward_hooks, with _is_full_backward_hook recording which kind you chose, and asking for both on one module raises rather than guessing.

The part that matters for checkpointing is what a hook is not. It is not in state_dict() and never was; deepcopy carries hooks along, and a freshly constructed module loaded from a checkpoint has none. So a debugging hook you registered in a notebook survives every copy you make in that session and vanishes the moment the model is rebuilt in a training script, silently, with the model still running.

import copy
import torch
from torch import nn

m = nn.Linear(2, 2)
log = []
m.register_forward_pre_hook(lambda mod, args: log.append("pre"))
m.register_forward_hook(lambda mod, args, out: log.append("fwd"))
m.register_full_backward_hook(lambda mod, gi, go: log.append("bwd"))
m(torch.randn(1, 2, requires_grad=True)).sum().backward()

print(log)                                                              # ['pre', 'fwd', 'bwd']
print(len(m._forward_hooks), len(m._backward_hooks), m._is_full_backward_hook)   # 1 1 True
print(list(m.state_dict()))                                             # ['weight', 'bias']
print(len(copy.deepcopy(m)._forward_hooks))                             # 1

fresh = nn.Linear(2, 2)
fresh.load_state_dict(m.state_dict())
print(len(fresh._forward_hooks))                                        # 0