autocast noise to cuda in ModelLoader

lucidrains · Nov 11, 2020 · e8a7dc8 · e8a7dc8
1 parent 440c68e
commit e8a7dc8
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/stylegan2_pytorch/stylegan2_pytorch.py b/stylegan2_pytorch/stylegan2_pytorch.py
@@ -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)

diff --git a/stylegan2_pytorch/version.py b/stylegan2_pytorch/version.py
@@ -1 +1 @@
-__version__ = '1.4.0'
+__version__ = '1.4.1'