use device from latents

lucidrains · Jan 13, 2021 · 8643aab · 8643aab
1 parent 7c16bed
commit 8643aab
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/lightweight_gan/lightweight_gan.py b/lightweight_gan/lightweight_gan.py
@@ -1157,7 +1157,7 @@ def calculate_fid(self, num_batches):
             for j in range(generated_images.size(0)):
                 torchvision.utils.save_image(generated_images[j, :, :, :], str(Path(fake_path) / f'{str(j + batch_num * self.batch_size)}-ema.{ext}'))
 
-        return fid_score.calculate_fid_given_paths([real_path, fake_path], 256, True, 2048)
+        return fid_score.calculate_fid_given_paths([real_path, fake_path], 256, latents.device, 2048)
 
     @torch.no_grad()
     def generate_truncated(self, G, style, trunc_psi = 0.75, num_image_tiles = 8):

diff --git a/lightweight_gan/version.py b/lightweight_gan/version.py
@@ -1 +1 @@
-__version__ = '0.16.2'
+__version__ = '0.16.3'