wrap code in ``` to make it easier to read.
why dont you simply assign a new layer in the same for loop? It depends on the type of model. For sequential you can directly change like a list, check the code below. If you are using vision models you can use .features to accsss each layer and reassign it with a new one
import torch

model = torch.nn.Sequential()
_ = model.add_module('a', torch.nn.Linear(1,2))
_ = model.add_module('b', torch.nn.Linear(3,4))
model[0] = torch.nn.Linear(5,6)
model

output: Sequential(
  (a): Linear(in_features=5, out_features=6, bias=True)
  (b): Linear(in_features=3, out_features=4, bias=True)
)