Fix pickling NeptuneLogger in offline mode

Lightning-AI · May 7, 2020 · a3d1fd4 · a3d1fd4
1 parent 8d0236f
commit a3d1fd4
Showing 1 changed file with 7 additions and 2 deletions.
diff --git a/pytorch_lightning/loggers/neptune.py b/pytorch_lightning/loggers/neptune.py
@@ -197,8 +197,12 @@ def __init__(self,
 
     def __getstate__(self):
         state = self.__dict__.copy()
-        # cannot be pickled
+
+        # Experiment cannot be pickled, and additionally its ID cannot be pickled in offline mode
         state['_experiment'] = None
+        if self.offline_mode:
+            state['_experiment_id'] = None
+
         return state
 
     @property
@@ -358,7 +362,8 @@ def _create_or_get_experiment(self):
         if self.offline_mode:
             project = neptune.Session(backend=neptune.OfflineBackend()).get_project('dry-run/project')
         else:
-            project = neptune.Session.with_default_backend(api_token=self.api_key).get_project(self.project_name)
+            session = neptune.Session.with_default_backend(api_token=self.api_key)
+            project = session.get_project(self.project_name)
 
         if self._experiment_id is None:
             exp = project.create_experiment(