for param in model.parameters():
    param.requires_grad = False

torch.manual_seed(123)
model.out_head = torch.nn.Linear(BASE_CONFIG["emb_dim"], 2)

for layer in (model.trf_blocks[-1], model.final_norm):
    for param in layer.parameters():
        param.requires_grad = True

        __
        __