I define a class with model initializer. When I wrap one model with that class
it looks fine, but with two model instances, it does not converge.
Code example:
def wrapper(nn.Module):

    def __init__(self, model):
          self.model = model
    def forward(self, x):
         return self.model(x)



model1 = net1()

w_model1 = wrapper(mode1)


model2 = net2()

#when I include this line, w_model1 does not converge.
w_model2 = wrapper(model2)