Fixing TPU training by disabling wandb.watch gradients logging for TPU (#4926)

bb154ac5 · Saurabh Misra · GitHub · fb6cccb8 · bb154ac5
Unverified Commit bb154ac5 authored Jun 17, 2020 by Saurabh Misra Committed by GitHub Jun 17, 2020
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 2 deletions

src/transformers/trainer.py src/transformers/trainer.py +2 -2

No files found.
--- a/src/transformers/trainer.py
+++ b/src/transformers/trainer.py
@@ -335,8 +335,8 @@ class Trainer:
                'Automatic Weights & Biases logging enabled, to disable set os.environ["WANDB_DISABLED"] = "true"'
            )
            wandb.init(project=os.getenv("WANDB_PROJECT", "huggingface"), config=vars(self.args))
-            # keep track of model topology and gradients
+            # keep track of model topology and gradients, unsupported on TPU
-            if os.getenv("WANDB_WATCH") != "false":
+            if not is_torch_tpu_available() and os.getenv("WANDB_WATCH") != "false":
                wandb.watch(
                    self.model, log=os.getenv("WANDB_WATCH", "gradients"), log_freq=max(100, self.args.logging_steps)
                )