fix: use ref only when necessaru

jolibrain · Jul 5, 2023 · ef2b0f2 · ef2b0f2
1 parent 13bd2d3
commit ef2b0f2
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/models/palette_model.py b/models/palette_model.py
@@ -535,7 +535,13 @@ def compute_palette_loss(self):
                 # the highest class is the unconditionned one.
                 cls = torch.where(drop_ids, self.num_classes - 1, cls)
 
-        ref = self.ref_A
+        if (
+            self.opt.alg_palette_cond_image_creation == "ref"
+            or "ref" in self.opt.alg_palette_conditioning
+        ):
+            ref = self.ref_A
+        else:
+            ref = None
 
         noise, noise_hat = self.netG_A(
             y_0=y_0, y_cond=y_cond, noise=noise, mask=mask, cls=cls, ref=ref