From 0fa9cf1c09db69591beec2b81fa2c56f93fd7850 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Florian=20M=C3=BCller-Fouarge?=
 <57919375+FlorianMF@users.noreply.github.com>
Date: Wed, 3 Mar 2021 16:11:30 +0000
Subject: [PATCH] Update variational.py

---
 pl_bolts/callbacks/variational.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pl_bolts/callbacks/variational.py b/pl_bolts/callbacks/variational.py
index 189b00f800..a6d7f7f04e 100644
--- 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)