Skip to content

← Devices and Data Movement step 4 of 4

Medium Primitives

Register a buffer that follows module conversions

Jump to the task

A layer may need a tensor that is not a learned weight, such as a scale or a mask. Assigning an ordinary tensor to self does not register it with the module. A dtype conversion makes the difference visible without a GPU:

import torch
import torch.nn as nn

class Layer(nn.Module):
    def __init__(self, n):
        super().__init__()
        self.w = nn.Parameter(torch.ones(n, dtype=torch.float32))
        self.scale = torch.full((n,), 2.0, dtype=torch.float32)

    def forward(self, x):
        return x * self.w * self.scale

layer = Layer(3)
layer.to(torch.float64)
print(layer.w.dtype)      # torch.float64
print(layer.scale.dtype)  # torch.float32

The weight was converted; scale stayed at its original dtype. Device moves follow the same registration rule. If you move this layer and its input to CUDA, the CPU scale can cause a device mismatch during forward.

Register state that belongs to the model

Replace the scale assignment in __init__ with:

# Inside Layer.__init__, after super().__init__():
self.register_buffer("scale", torch.full((n,), 2.0))

Module.to() visits parameters and buffers in the module and its registered children. It moves them to the requested device and converts floating-point or complex tensors to the requested dtype. Integer and boolean buffers keep their dtype when you request a floating-point cast.

A buffer is excluded from parameters(), so an optimizer built from that iterator will not receive it. Registering a buffer does not itself disable gradients; the scale above has no gradients because torch.full creates it that way by default.

Decide what belongs in a checkpoint

By default, buffers appear in state_dict() alongside parameters. This matters for state that cannot simply be recreated, such as accumulated running statistics. A plain tensor attribute is omitted unless you implement custom saving.

For a tensor you can reliably rebuild from the model configuration, use register_buffer(name, tensor, persistent=False) to leave it out of state_dict(). It still follows module device moves and applicable dtype conversions. Use this deliberately: the checkpoint will not restore its value.

Your task

Fix Layer so scale is a persistent buffer initialized to a vector of twos. Keep w as a Parameter initialized to ones. Then implement:

def scale_dtype_after_cast(n: int) -> dict

Build Layer(n), call layer.to(torch.float64), and return these three fields:

  • w: the weight's dtype name, without the torch. prefix.
  • scale: the scale's dtype name, without the torch. prefix.
  • in_state_dict: whether scale appears in the module's state_dict().

For n=3, return {"w": "float64", "scale": "float64", "in_state_dict": True}. The checks use a CPU dtype conversion, so no CUDA device is needed.