nn.Sequential runs layers one after another. A model that does anything else, such as adding its input back to its output, is a class of your own: subclass nn.Module, create the parts in __init__, and use them in forward.

import torch
from torch import nn

class Affine(nn.Module):
    def __init__(self, size):
        super().__init__()
        self.scale = nn.Parameter(
            torch.ones(size))
        self.shift = nn.Parameter(
            torch.zeros(size))

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

m = Affine(3)
for name, p in m.named_parameters():
    print(name, p.shape)
print(m(torch.ones(2, 3)))

Wrapping a tensor in nn.Parameter and assigning it to self registers it: it appears in parameters(), so the optimizer trains it, and in state_dict(), so it is saved. A layer assigned to self registers its own parameters the same way. A plain tensor stays invisible to both.

Layers start from random values scaled to their size. nn.init replaces them in place, for example nn.init.zeros_(layer.bias) or nn.init.xavier_uniform_(layer.weight).

Write ReZero(nn.Module) taking size: it holds self.linear = nn.Linear(size, size), with the weight set by nn.init.xavier_uniform_ and the bias by nn.init.zeros_, and a learnable self.alpha of shape (1,) starting at 0. Its forward(x) returns x + self.alpha * torch.tanh(self.linear(x)).

The explanation opens here once every check passes.