Initialise the whole tree

Custom initialisation has to reach every layer, including layers inside layers inside a Sequential you built from a list.

def init(module):
    if isinstance(module, nn.Linear):
        nn.init.zeros_(module.weight)
        nn.init.zeros_(module.bias)

model.apply(init)

apply walks the module tree depth-first and calls the function on every module, then on the root. It is the same registry walk that parameters(), .to() and train() use, which is why registering your submodules properly is what makes it work.

Why not a loop over parameters

for p in model.parameters():
    nn.init.zeros_(p)

loses the information you need. A Linear‘s weight and bias want different treatment, and so do a LayerNorm‘s: the standard recipe initialises weights from a scaled normal, biases to zero, and normalisation weights to one. parameters() gives you tensors with no idea what they are for.

apply gives you the module, so isinstance can decide.

no_grad, and why you rarely write it here

nn.init.* functions are decorated with torch.no_grad() internally, so they can write to a leaf that requires grad without the error from the memory track. Hand-written initialisation is not:

with torch.no_grad():
    module.weight.mul_(0.5)        # needs the block
nn.init.zeros_(module.weight)      # does not

Knowing which is which stops the error being mysterious when it appears.

The related one-liner

sum(p.numel() for p in model.parameters())

The parameter count everybody quotes. Note it counts tensors’ elements, deduplicated by parameters(), so tied weights are counted once, which is the honest number.

Your task

def init_and_count(depth: int) -> dict

Build a model with depth Linear(2, 2) layers inside a Sequential, set every Linear weight to all ones and every bias to zero using apply, and return the total of all parameter values, the total number of parameter elements, and the output of a forward pass on [1.0, 1.0].