change animegan2 params

Guo-HY · Dec 8, 2022 · 98426b4 · 98426b4
1 parent 3f33802
commit 98426b4
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/src/animegan2/test.py b/src/animegan2/test.py
@@ -38,11 +38,11 @@ def to_32s(x):
 
 def test(output_dir, input_dir, image_name):
     device = 'cpu'
-
+    PT = './animegan2/weights/celeba_distill.pt'
     net = Generator()
-    net.load_state_dict(torch.load('./animegan2/weights/paprika.pt', map_location="cpu"))
+    net.load_state_dict(torch.load(PT, map_location="cpu"))
     net.to(device).eval()
-    print(f"model loaded: {'./animegan2/weights/paprika.pt'}")
+    print(f"model loaded: {PT}")
 
     os.makedirs(output_dir, exist_ok=True)
 

diff --git a/src/media/5.jpg b/src/media/5.jpg