Skip to content

A
animegan2-pytorch
Commit b8e3a27d authored by xhlulu's avatar xhlulu
Browse files
Options

Fix return function

parent 1799501f
Hide whitespace changes
Inline Side-by-side
Showing with 5 additions and 2 deletions
hubconf.py
View file @ b8e3a27d
import torch import torch
from model import Generator
def animegan2(pretrained=True, device="cpu", progress=True, check_hash=True): def animegan2(pretrained=True, device="cpu", progress=True, check_hash=True):
from model import Generator
model = Generator() model = Generator()
if type(pretrained) == str: if type(pretrained) == str:
ckpt_url = pretrained ckpt_url = pretrained
pretrained = True pretrained = True
...@@ -18,3 +19,5 @@ def animegan2(pretrained=True, device="cpu", progress=True, check_hash=True): ...@@ -18,3 +19,5 @@ def animegan2(pretrained=True, device="cpu", progress=True, check_hash=True):
check_hash=check_hash, check_hash=check_hash,
) )
model.load_state_dict(state_dict) model.load_state_dict(state_dict)
return model
\ No newline at end of file
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment