enable log_code in wandb callback

michalwols · Mar 23, 2024 · 96f5b0c · 96f5b0c
1 parent 8932acd
commit 96f5b0c
Showing 1 changed file with 9 additions and 1 deletion.
diff --git a/yann/callbacks/wandb.py b/yann/callbacks/wandb.py
@@ -19,7 +19,7 @@ def __init__(
       entity=None,
       name=None,
       watch_freq=0,
-      # log_code=True,
+      log_code=True,
       batch_log_freq=10,
       trackers=None,
   ):
@@ -30,6 +30,7 @@ def __init__(
     self.name = name
     self.batch_log_freq = batch_log_freq
     self.watch_freq = watch_freq
+    self.log_code = log_code
 
     self.trainer = None
 
@@ -48,6 +49,13 @@ def on_train_start(self, trainer: 'yann.train.Trainer' = None):
         config=dict(trainer.params) if trainer.params else {}
       )
 
+      if self.log_code is True:
+        self.run.log_code('.')
+      elif isinstance(self.log_code, dict):
+        self.run.log_code(**self.log_code)
+      elif isinstance(self.log_code, str):
+        self.run.log_code(self.log_code)
+
     if trainer.model and self.watch_freq:
       self.run.watch(
         models=trainer.model,