Hey @gphilip , seems like I’ve fixed my issues. This is my final solution:
shift_layer = generator.main[11]
optimizer_b = optim.SGD(shift_layer.parameters(), lr = ...)

and in the training loop I just call optimizer_b.step() for every epoch.
So far this seems to be working! Thanks to you both for your help. Appreciate it!