Hi all,
Coming from Keras (TensorFlow), I used to have a somewhat robust way to save/load models trained in previous development iterations. Now using PyTorch, I learned that it’s preferable to save only the model parameters contained in the model.state_dict(). I would like to be able to edit my model code, retrain the model and then perform various experiments on this last model and also its predecessors. I can only think of two naive approaches:

Have a model definition function that supports a version argument. Save the model version with the model state dict and, during loading, use it to build the good model.

pros: explicit
cons: would get messy pretty quickly, would require lots of discipline to keep the old code unchanged

Save the code of the model with the model state dict. During loading, dynamically use this code to build the model.

pros: less messy
cons: would probably get complex if the model definition function and its dependencies don’t fit in a single file (which, excluding libraries, is not my case for now)
Cleary both approaches could still break if I change my python environment (update the project dependencies) but I guess I can live with that. Do you think these approaches are viable? Do you know better?