antonypotapenko/cyclegan-lab5

1from __future__ import annotations2 3from pathlib import Path4 5import torch6 7from modeling import load_generator8 9 10ARTIFACTS = Path("artifacts")11INPUT_SHAPE = (1, 3, 256, 256)12 13 14def export_one(src_name: str, dst_name: str) -> None:15    src = ARTIFACTS / src_name16    dst = ARTIFACTS / dst_name17    model = load_generator(src, device="cpu")18    example = torch.randn(*INPUT_SHAPE)19    traced = torch.jit.trace(model, example)20    traced.save(str(dst))21    print(f"Saved {dst}")22 23 24if __name__ == "__main__":25    export_one("gen_a2b_custom.pth", "gen_a2b_custom_jit.pt")26    export_one("gen_b2a_custom.pth", "gen_b2a_custom_jit.pt")27