Set to zero, not None

huggingface · Sep 12, 2023 · 60ea398 · 60ea398
1 parent 535ad35
commit 60ea398
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
@@ -2021,7 +2021,7 @@ def check_trigger(self):
             self.flag_tensor = torch.tensor(0, device=self.device)
         flag_tensor = self.reduce(self.flag_tensor)
         if flag_tensor.item() >= 1:
-            self.flag_tensor = None
+            self.flag_tensor = torch.tensor(0, device=self.device)
             return True
         return False