import torch def save_model(model): torch.save(model.state_dict(), 'model_weights.pth') def load_model(model): return model.load_state_dict(torch.load('./models/model_weights_27_styles.pth', map_location=torch.device('cpu')))