From fd2174664c8879c747ada3e6e0a2486858808421 Mon Sep 17 00:00:00 2001
From: Julien Chaumond <chaumond@gmail.com>
Date: Tue, 5 May 2020 10:59:23 -0400
Subject: [PATCH] [Trainer] W&B: Enable model watch

See https://github.com/huggingface/transformers/pull/3916
---
 src/transformers/trainer.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py
index d00cabc072..1f3c7cf24a 100644
--- a/src/transformers/trainer.py
+++ b/src/transformers/trainer.py
@@ -226,10 +226,14 @@ class Trainer:
         return optimizer, scheduler
 
     def _setup_wandb(self):
-        # Start a wandb run and log config parameters
+        """
+        Setup the optional Weights & Biases (`wandb`) integration.
+
+        One can override this method to customize the setup if needed.
+        """
         wandb.init(name=self.args.logging_dir, config=vars(self.args))
         # keep track of model topology and gradients
-        # wandb.watch(self.model)
+        wandb.watch(self.model)
 
     def train(self, model_path: Optional[str] = None):
         """