import sys
import torch
from omegaconf import OmegaConf

from engine.util.build import build_model

config = OmegaConf.load(sys.argv[1])


generator = build_model(config.model.generator)

ckp = torch.load(sys.argv[2], map_location="cpu")

generator.module.load_state_dict(ckp["generator_main"])