15 lines
295 B
Python
15 lines
295 B
Python
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"])
|