After the backward call, and before the optimizier update, you can access the gradient tensor (param.grad) and assign zero to the parts that you want to freeze/not update. This way these weights are left unchanged after the optim.step call.