From 4625b34f4ebcd30ffb24f1b03ece8a8362c7bfb2 Mon Sep 17 00:00:00 2001 From: Kohya S Date: Tue, 29 Apr 2025 21:27:04 +0900 Subject: [PATCH] Fix mean image aspect ratio error calculation to avoid NaN values --- library/train_util.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/library/train_util.py b/library/train_util.py index 100ef475..b9d08f25 100644 --- a/library/train_util.py +++ b/library/train_util.py @@ -957,8 +957,11 @@ class BaseDataset(torch.utils.data.Dataset): self.bucket_info["buckets"][i] = {"resolution": reso, "count": len(bucket)} logger.info(f"bucket {i}: resolution {reso}, count: {len(bucket)}") - img_ar_errors = np.array(img_ar_errors) - mean_img_ar_error = np.mean(np.abs(img_ar_errors)) + if len(img_ar_errors) == 0: + mean_img_ar_error = 0 # avoid NaN + else: + img_ar_errors = np.array(img_ar_errors) + mean_img_ar_error = np.mean(np.abs(img_ar_errors)) self.bucket_info["mean_img_ar_error"] = mean_img_ar_error logger.info(f"mean ar error (without repeats): {mean_img_ar_error}")