Update cleanrl/ppo_atari_accelerate.py

Co-authored-by: Costa Huang <[email protected]>

cleanrl/ppo_atari_accelerate.py

@@ -243,8 +243,10 @@ def get_action_and_value(self, x, action=None):
             rewards[step] = torch.tensor(reward).to(device).view(-1)
             next_obs, next_done = torch.Tensor(next_obs).to(device), torch.Tensor(next_done).to(device)
 
-
-            if "final_info" in infos and accelerator.is_main_process:
+            if not writer:
+                continue
+
+            if "final_info" in infos:
                 for info in infos["final_info"]:
                     if info and "episode" in info:
                         print(f"global_step={global_step}, episodic_return={info['episode']['r']}")