Is it possible to instantiate a model from a saved model?

1

Typically, I use the following code to resume my model. Now, I am planning to use multi models on the same code. I have models which are independently implemented, and I also have the each saved model after trainings.
Then, I would like to instantiate each model. Though, I need arguments for each model. It is not difficult, but feel like a kind of redundant.

def save_checkpoint(state, filename):
    torch.save(state, filename)

def save(..):
    save_checkpoint({
                'epoch'      : epoch,
                'state_dict' : model.state_dict(),
                'optimizer'  : optimizer.state_dict()
    }, filename=filename)


model = MyModel(**args)

ckpt = torch.load(args.resume)
model.load_state_dict(ckpt['state_dict'])
optimizer.load_state_dict(ckpt['optimizer'])

The following code is my ideal since I only set the saved models’ path. Is it possible? I think I need to modify my save function.

# model1 = MyModel1(a, b, c, d, e) # ugh, i want to skip this 
model1 = load_and_instantiate_model(args.resume1)
model2 = load_and_instantiate_model(args.resume2)
model2 = load_and_instantiate_model(args.resume3)
2

One solution may be storing all arguments.

def save(..):
    save_checkpoint({
                'kargs'      : kargs,
                'epoch'      : epoch,
                'state_dict' : model.state_dict(),
                'optimizer'  : optimizer.state_dict()
    }, filename=filename)

kargs = {...}
model = MyModel(**kargs)
##--------------------##
ckpt = torch.load(args.resume)
model = MyModel(**ckpt[kargs])