From 7218588c03be1e7fa5566b8836826e6b1c9065d2 Mon Sep 17 00:00:00 2001 From: Volpeon Date: Sat, 24 Dec 2022 08:32:02 +0100 Subject: Fix --- training/util.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) (limited to 'training') 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 -- cgit v1.2.3-70-g09d2