Update variational.py

Lightning-Universe · Mar 3, 2021 · 0fa9cf1 · 0fa9cf1
1 parent 6ee3a0b
commit 0fa9cf1
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/pl_bolts/callbacks/variational.py b/pl_bolts/callbacks/variational.py
@@ -67,7 +67,7 @@ def on_epoch_end(self, trainer: Trainer, pl_module: LightningModule) -> None:
             )
             images = torch.cat(images, dim=0)  # type: ignore[assignment]
 
-            num_rows = int(math.sqrt(self.steps))
+            num_rows = self.steps
             grid = torchvision.utils.make_grid(images, nrow=num_rows, normalize=self.normalize)
             str_title = f'{pl_module.__class__.__name__}_latent_space'
             trainer.logger.experiment.add_image(str_title, grid, global_step=trainer.global_step)