condition and supervise

broadinstitute · Jan 17, 2025 · 21039b1 · 21039b1
1 parent 5ffb908
commit 21039b1
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/ml4h/models/diffusion_blocks.py b/ml4h/models/diffusion_blocks.py
@@ -877,7 +877,7 @@ def test_step(self, batch):
         # this is computationally demanding, kid_diffusion_steps has to be small
         if self.input_map.axes() == 3 and self.inspect_model:
             images = self.denormalize(images)
-            generated_images = self.generate(
+            generated_images = self.generate(control_embed,
                 num_images=self.batch_size, diffusion_steps=20
             )
             self.kid.update_state(images, generated_images)