Thanks for the tip, I think that works for me. It seems some care needs to be taken not to set this on leaf parameters
layer1 = nn.Linear(1,1, bias=False)
layer2 = nn.Linear(1,1, bias=False)
net = nn.Sequential(layer1, layer2)
x = torch.ones((1))
layer2.weight.requires_grad=False
net(x*x).backward()
assert layer2.weight.grad is None
layer1.weight.requires_grad=False
net(x*x).backward()   # nothing to propagate