What a move does

Every conversion a module can perform goes through _apply: .to(), .cuda(), .float(), .half(), to_empty(). It recurses into children, then walks _parameters and _buffers, applies the function under no_grad, and returns self. That last word is the difference from tensors: t.to(...) hands you a new tensor and leaves t alone, while m.to(...) rewrites m and hands it back.

Inside the parameter loop there is a fork, and it decides more than it looks like it does. compute_should_use_set_data asks torch._has_compatible_shallow_copy_type(tensor, tensor_applied), and on a yes it writes param.data = param_applied, keeping the Parameter object and giving it different storage. On a no it constructs a new Parameter and puts that in _parameters instead.

A dtype change on the same device takes the first branch. A device change takes the second. Everything in this lesson follows from which branch ran.

        for key, param in self._parameters.items():
            if param is None:
                continue
            # Tensors stored in modules are graph leaves, and we don't want to
            # track autograd history of `param_applied`, so we have to use
            # `with torch.no_grad():`
            with torch.no_grad():
                param_applied = fn(param)
            should_use_set_data = compute_should_use_set_data(param, param_applied)
            if should_use_set_data:
                param.data = param_applied
                out_param = param
            else:
                assert isinstance(param, Parameter)
                assert param.is_leaf
                out_param = Parameter(param_applied, param.requires_grad)
                self._parameters[key] = out_param

Ask for the dtype a module already has and nothing happens at all: t.to(torch.float32) on a float32 tensor returns the same tensor, so param.data is reassigned to itself and the storage address never changes. That is why calling .float() defensively costs nothing.

Ask for a different dtype and the object identity survives while the storage does not. The measured result is the same Parameter at a new address, with requires_grad carried over by the set_data path and grads converted alongside in the loop just below.

So anything holding a reference to a parameter keeps working across a dtype change. That includes the optimizer, which is the reason the next failure is confusing when you meet it.

import torch
from torch import nn

m = nn.Linear(4, 4)
p, ptr = m.weight, m.weight.data_ptr()

m.to(torch.float32)                       # already float32
print(m.weight is p, m.weight.data_ptr() == ptr)          # True True

m.to(torch.float64)                       # a real conversion, same device
print(m.weight is p, m.weight.data_ptr() == ptr,
      m.weight.requires_grad)                             # True False True

n = nn.Linear(4, 4)
q = n.weight
n.to("meta")                              # a device change
print(n.weight is q,
      torch._has_compatible_shallow_copy_type(q, n.weight))   # False False

Cross a device boundary and _has_compatible_shallow_copy_type says no, so _apply installs brand new Parameter objects in _parameters. The old ones are still alive, still holding their storage, and still sitting in whatever list you handed them to earlier.

An optimizer built before the move is still holding the originals. The arc under chapter 4 establishes why: model.parameters() is drained into a flat list of tensor references at construction, and nothing re-reads it. So the optimizer keeps stepping, it raises nothing, and it updates tensors the module no longer refers to. The run below does it on this machine using the meta device, which is the second device available here: after the move, the optimizer's parameter is the old object, that old object has moved, and the module's parameter has not.

This is the mechanism under the usual advice to construct the optimizer after moving the model. The advice is right; the reason is not that PyTorch loses track of your parameters, but that a device move is defined as replacement rather than mutation.

A dtype change edits your parameters. A device change replaces them, and every reference you kept still points at the originals.

import torch
from torch import nn

torch.manual_seed(0)
m = nn.Linear(4, 4)
opt = torch.optim.SGD(m.parameters(), lr=1.0)
m(torch.randn(8, 4)).sum().backward()

orphan = m.weight
before = orphan.detach().clone()
m.to("meta")                               # a device change: new Parameter objects
opt.step()

print(opt.param_groups[0]["params"][0] is orphan)   # True
print(torch.equal(orphan.detach(), before))         # False: the step moved the orphan
print(m.weight.is_meta, m.weight.data_ptr())        # True 0

The dtype case fails differently, and later. Adam's per-parameter state is allocated on that parameter's first step and in that parameter's dtype, which the arc under chapter 4 measures from the optimizer side. _apply never touches the optimizer, so convert the model afterwards and the parameter is float64 while exp_avg is still float32.

The lookup still works, because the state is keyed by the Parameter object and that object survived. The next opt.step() is where it lands: RuntimeError: expected dtype float for end but got dtype double, raised from exp_avg.lerp_(grad, 1 - beta1) inside Adam.

Ordering is what removes the whole class. Decide device and dtype before the optimizer exists, and treat a mid-run conversion as a rebuild: new optimizer, then reload its state dict if the run has to continue from where it was.

import torch
from torch import nn

torch.manual_seed(0)
m = nn.Linear(4, 4)
opt = torch.optim.Adam(m.parameters(), lr=1e-2)
m(torch.randn(8, 4)).sum().backward()
opt.step()
print(opt.state[m.weight]["exp_avg"].dtype)          # torch.float32

m.to(torch.float64)                        # same device, so the objects survive
print(m.weight.dtype, m.weight in opt.state,
      opt.state[m.weight]["exp_avg"].dtype)          # torch.float64 True torch.float32

opt.zero_grad()
m(torch.randn(8, 4).double()).sum().backward()
opt.step()
# RuntimeError: expected dtype float for `end` but got dtype double

Module.to does not hand _apply your arguments directly. It builds a closure called convert at module.py:1146, and the decisive expression inside it reads dtype if t.is_floating_point() or t.is_complex() else None. The dtype reaches a tensor only when that test passes, so an integer buffer keeps its own dtype while the parameters around it change.

That rule is what keeps a step counter, a token id table or a boolean mask intact through model.half(). It also means a dtype argument is not a promise about every tensor in the module, which matters when you are reasoning about what a mixed-dtype forward pass will do.

The plain attribute from lesson one is skipped for the older reason: _apply walks two registries, and it was never in either. A module converted to float64 can therefore be carrying an int64 buffer and a float32 tensor on self at the same time, all three correct by the rules and only one of them what you meant.

import torch
from torch import nn


class Step(nn.Module):
    def __init__(self):
        super().__init__()
        self.lin = nn.Linear(2, 2)
        self.register_buffer("count", torch.zeros(1, dtype=torch.int64))
        self.scale = torch.ones(1)              # a plain attribute


m = Step().to(torch.float64)
print(m.lin.weight.dtype, m.count.dtype, m.scale.dtype)
# torch.float64 torch.int64 torch.float32

One device allocates nothing. Build a module under torch.device('meta') and its parameters have real shapes and dtypes with a data_ptr of zero, which is how you get the structure of a model too large to allocate twice, or how you inspect shapes without paying for them.

From there, two ways forward. to_empty(device='cpu') allocates storage of the right shape and leaves it uninitialized, ready for a normal load; the values in between are whatever the allocator handed back, so reading them before the load is a bug. Or load with assign=True, which adopts the checkpoint's tensors and allocates nothing of its own.

The path that does not work is the ordinary copy load onto a meta module. copy_ into a tensor with no storage writes nowhere, so the load warns per key that it is a no-op and suggests assign=True, and the module stays meta. Read the warning rather than the missing exception; there is no error here, only a model that never got its weights.

The same storage question runs through initialization. reset_parameters writes into the storage the module already has, in place, so it keeps the address and the parameter identity, which is why re-initializing a model does not invalidate an optimizer while moving it does.

import torch
from torch import nn

torch.manual_seed(0)
trained = nn.Linear(4, 4)

skeleton = nn.Linear(4, 4, device="meta")
print(skeleton.weight.device, skeleton.weight.data_ptr())   # meta 0

a = nn.Linear(4, 4, device="meta")
a.to_empty(device="cpu")                   # allocate, uninitialized
a.load_state_dict(trained.state_dict())
print(a.weight.data_ptr() != 0, torch.equal(a.weight, trained.weight))   # True True

ck = trained.state_dict()
b = nn.Linear(4, 4, device="meta")
b.load_state_dict(ck, assign=True)         # adopt the checkpoint's tensors
print(b.weight.device, b.weight.data_ptr() == ck["weight"].data_ptr())   # cpu True

c = nn.Linear(4, 4, device="meta")
c.load_state_dict(trained.state_dict())    # the copy path, onto no storage
# UserWarning: for weight: copying from a non-meta parameter in the checkpoint
# to a meta parameter in the current model, which is a no-op.
print(c.weight.is_meta)                    # True