utils/save_load.py

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')))