Fix return function

safarri · Nov 6, 2021 · b8e3a27 · b8e3a27
1 parent 1799501
commit b8e3a27
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/hubconf.py b/hubconf.py
@@ -1,9 +1,10 @@
 import torch
-from model import Generator
-
 
 def animegan2(pretrained=True, device="cpu", progress=True, check_hash=True):
+    from model import Generator
+
     model = Generator()
+
     if type(pretrained) == str:
         ckpt_url = pretrained
         pretrained = True
@@ -18,3 +19,5 @@ def animegan2(pretrained=True, device="cpu", progress=True, check_hash=True):
             check_hash=check_hash,
         )
         model.load_state_dict(state_dict)
+
+    return model