Thank you for your response but this doesn’t work for me. This passes my test_model_change function so the weights do get reinitialized but when I try to retrain my model, the loss does not decrease. Interestingly, when I reset the model to the original weight parameters at the first initialization I don’t face this issue, only when I randomly try to set weights to some values.
def reset_weights_rand(model):
    rand_state_dict = pickle.load(open('rand_dict.pkl', 'rb'))
    for name, param in model.named_parameters():
        param.data = rand_state_dict[name]
rand_state_dict = copy.deepcopy(initial_state_dict)

for key,value in rand_state_dict.items():
    if 'bias' in key:
        init.normal_(value.data)
    else:
        init.xavier_normal_(value.data)        

with open(r"rand_dict.pkl", "wb") as output_file:
    pickle.dump(rand_state_dict, output_file)

train(model,1,train_loader,device,optimizer) # train for one epoch to test

full_cap = copy.deepcopy(model.state_dict())

reset_weights_rand(model)

# model.apply(reinit_weights) # this reinitalizes the weight parameter but doesn't change the new_weight parameter
test_model_change(full_cap, model) # passes now

# train 90% capacity for 10 more epochs
optimizer2 = torch.optim.Adam(model.parameters(), lr=lr) # problematic

train(model,5, train_loader,device,optimizer2)


image1530×1480 342 KB

As you can see the loss does not decrease even though no pruning as been done.