Skip to content

Commit

Permalink
autocast noise to cuda in ModelLoader
Browse files Browse the repository at this point in the history
  • Loading branch information
lucidrains committed Nov 11, 2020
1 parent 440c68e commit e8a7dc8
Show file tree
Hide file tree
Showing 2 changed files with 2 additions and 1 deletion.
1 change: 1 addition & 0 deletions stylegan2_pytorch/stylegan2_pytorch.py
Original file line number Diff line number Diff line change
Expand Up @@ -1261,6 +1261,7 @@ def __init__(self, *, base_dir, name = 'default', load_from = -1):
self.model.load(load_from)

def noise_to_styles(self, noise, trunc_psi = None):
noise = noise.cuda()
w = self.model.GAN.S(noise)
if exists(trunc_psi):
w = self.model.truncate_style(w)
Expand Down
2 changes: 1 addition & 1 deletion stylegan2_pytorch/version.py
Original file line number Diff line number Diff line change
@@ -1 +1 @@
__version__ = '1.4.0'
__version__ = '1.4.1'

0 comments on commit e8a7dc8

Please sign in to comment.