import torch
from test_model import MyModelConfig,MyModel


MyModelConfig.register_for_auto_class()
MyModel.register_for_auto_class("AutoModel")
save_config = MyModelConfig(input_dim=10,layers_num=3)
save_config.save_pretrained("custom-mymodel")

mymodel = MyModel(save_config)
mymodel.load_state_dict(torch.load('pytorch_model.bin'))
mymodel.save_pretrained("custom-mymodel")

mymodel.push_to_hub("custom-mymodel_v1")