Skip to content

Commit 4625b34

Browse files
committed
Fix mean image aspect ratio error calculation to avoid NaN values
1 parent 3b25de1 commit 4625b34

File tree

1 file changed

+5
-2
lines changed

1 file changed

+5
-2
lines changed

library/train_util.py

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -957,8 +957,11 @@ def make_buckets(self):
957957
self.bucket_info["buckets"][i] = {"resolution": reso, "count": len(bucket)}
958958
logger.info(f"bucket {i}: resolution {reso}, count: {len(bucket)}")
959959

960-
img_ar_errors = np.array(img_ar_errors)
961-
mean_img_ar_error = np.mean(np.abs(img_ar_errors))
960+
if len(img_ar_errors) == 0:
961+
mean_img_ar_error = 0 # avoid NaN
962+
else:
963+
img_ar_errors = np.array(img_ar_errors)
964+
mean_img_ar_error = np.mean(np.abs(img_ar_errors))
962965
self.bucket_info["mean_img_ar_error"] = mean_img_ar_error
963966
logger.info(f"mean ar error (without repeats): {mean_img_ar_error}")
964967

0 commit comments

Comments
 (0)