Add assertion for checking dataset size

Yuki-11 · Apr 29, 2024 · 62e5a2c · 62e5a2c
1 parent bb25483
commit 62e5a2c
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/model/engine/inference.py b/model/engine/inference.py
@@ -21,6 +21,7 @@ def inference_for_ss(args, cfg, model, test_loader):
 
  fnames = []
  max_iter = len(test_loader)
+ assert max_iter != 0, "Dataset size is 0!!"
 
  img_psnr_scores = np.array([])
  kernel_psnr_scores = np.array([])