Merge pull request #57 from neptune-ai/fix-lr-multi-gpu

fix learning_rate logging on multigpu
neptune-ai · Jul 24, 2023 · efdc4e4 · efdc4e4
2 parents 3b3dec4 + ac30d95
commit efdc4e4
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,3 +1,8 @@
+## neptune-tensorflow-keras 2.2.0
+
+### Fixes
+- Support logging learning_rate in Multi-GPU training set-up ([#57](https://github.com/neptune-ai/neptune-tensorflow-keras/pull/57))
+
 ## neptune-tensorflow-keras 2.1.0
 
 ### Changes

diff --git a/src/neptune_tensorflow_keras/impl/__init__.py b/src/neptune_tensorflow_keras/impl/__init__.py
@@ -164,7 +164,7 @@ def on_train_batch_end(self, batch, logs=None):
             self._log_metrics(logs, "train", "batch")
 
     def on_epoch_begin(self, epoch, logs=None):
-        self._model_logger["learning_rate"].append(self.model.optimizer.learning_rate)
+        self._model_logger["learning_rate"].append(self.model.optimizer.learning_rate.numpy())
 
     def on_epoch_end(self, epoch, logs=None):
         self._log_metrics(logs, "train", "epoch")