I am trying to implement a neuron pruning algorithm. Along training, during certain steps, i prune one neuron and that changes the architecture of my model (I basically recreate a new smaller network based on the bigger network). I need to save the model after each pruning. I am looking for a descent way to save the model.
The suggested way torch.save(model.state_dict(), PATH) does not work for me since to load the model later, i need to know the architecture and since i am changing the architecture constantly, i do not know the architecture for each model (unless there is a way to save the architecture separately which i do not know).
Also, torch.save(model., PATH) works, but problem is that i am doing a lot of stuff during training and for whatever reason the saved model using this method is huge (15 gb for a small model that saving it using torch.save(model.state_dict(), PATH)takes only 90 mb).
Is there any descent way to handle this situation?