From 8e3a33e0019228821e2c50f027e8db468a3d5c99 Mon Sep 17 00:00:00 2001
From: James Betker <jbetker@gmail.com>
Date: Wed, 16 Jun 2021 16:27:09 -0600
Subject: [PATCH] Fix a bug where non-rank-0 is computing FID before all images
 are saved.

---
 codes/trainer/eval/sr_diffusion_fid.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/codes/trainer/eval/sr_diffusion_fid.py b/codes/trainer/eval/sr_diffusion_fid.py
index 1891487f..74e834b5 100644
--- a/codes/trainer/eval/sr_diffusion_fid.py
+++ b/codes/trainer/eval/sr_diffusion_fid.py
@@ -55,5 +55,8 @@ class SrDiffusionFidEvaluator(evaluator.Evaluator):
                     torchvision.utils.save_image(g, osp.join(fid_fake_path, f"{counter}.png"))
                     counter += 1
 
-        return {"fid": fid_score.calculate_fid_given_paths([self.fid_real_samples, fid_fake_path], self.fid_batch_size,
-                                                           True, 2048)}
\ No newline at end of file
+        if self.env['rank'] <= 0:
+            return {"fid": fid_score.calculate_fid_given_paths([self.fid_real_samples, fid_fake_path], self.fid_batch_size,
+                                                               True, 2048)}
+        else:
+            return {}