make sure contrastive loss regularization and transparent images are …

…not both turned on
lucidrains · Jun 27, 2020 · c74e436 · c74e436
1 parent 7d0af67
commit c74e436
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/setup.py b/setup.py
@@ -4,7 +4,7 @@
   name = 'stylegan2_pytorch',
   packages = find_packages(),
   scripts=['bin/stylegan2_pytorch'],
-  version = '0.16.1',
+  version = '0.16.2',
   license='GPLv3+',
   description = 'StyleGan2 in Pytorch',
   author = 'Phil Wang',

diff --git a/stylegan2_pytorch/stylegan2_pytorch.py b/stylegan2_pytorch/stylegan2_pytorch.py
@@ -529,6 +529,7 @@ def __init__(self, image_size, latent_dim = 512, style_depth = 8, network_capaci
         self.GE = Generator(image_size, latent_dim, network_capacity, transparent = transparent, attn_layers = attn_layers, no_const = no_const)
 
         # experimental contrastive loss discriminator regularization
+        assert not (transparent and cl_reg), 'contrastive loss regularization does not work with transparent images yet'
         self.D_cl = ContrastiveLearner(self.D, image_size, hidden_layer='flatten') if cl_reg else None
 
         # wrapper for augmenting all images going into the discriminator