Fix return function

b8e3a27d · xhlulu · 1799501f · b8e3a27d
Commit b8e3a27d authored Nov 06, 2021 by xhlulu
Show whitespace changes
Inline Side-by-side

Showing with 5 additions and 2 deletions

hubconf.py hubconf.py +5 -2

No files found.
--- a/hubconf.py
+++ b/hubconf.py
 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
\ No newline at end of file