passing experiment to wandb (#720)

This commit is contained in:
Cole Hurwitz
2020-01-21 11:20:45 -05:00
committed by William Falcon
parent 9e654c4ec8
commit 707bcb2827
+3 -3
View File
@@ -33,7 +33,7 @@ class WandbLogger(LightningLoggerBase):
"""
def __init__(self, name=None, save_dir=None, offline=False, id=None, anonymous=False,
version=None, project=None, tags=None):
version=None, project=None, tags=None, experiment=None):
super().__init__()
self._name = name
self._save_dir = save_dir
@@ -41,7 +41,7 @@ class WandbLogger(LightningLoggerBase):
self._id = version or id
self._tags = tags
self._project = project
self._experiment = None
self._experiment = experiment
self._offline = offline
def __getstate__(self):
@@ -81,7 +81,7 @@ class WandbLogger(LightningLoggerBase):
@rank_zero_only
def log_metrics(self, metrics, step=None):
metrics["global_step"] = step
self.experiment.history.add(metrics)
self.experiment.log(metrics)
def save(self):
pass