1 files changed, 4 insertions, 1 deletions

diff --git a/training/util.py b/training/util.py
index 000173d..a623dc5 100644
--- a/training/util.py
+++ b/training/util.py
@@ -74,6 +74,9 @@ class CheckpointerBase:
             generator=generator,
         )
 
+        grid_cols = max(self.sample_batch_size, 4)
+        grid_rows = self.sample_batches * self.sample_batch_size / grid_cols
+
         with torch.autocast("cuda"), torch.inference_mode():
             for pool, data, latents in [("stable", val_data, stable_latents), ("val", val_data, None), ("train", train_data, None)]:
                 all_samples = []
@@ -119,7 +122,7 @@ class CheckpointerBase:
 
                     del samples
 
-                image_grid = make_grid(all_samples, self.sample_batches, self.sample_batch_size)
+                image_grid = make_grid(all_samples, grid_rows, grid_cols)
                 image_grid.save(file_path, quality=85)
 
                 del all_samples