feat(trainer): save checkpoints when load no data (#1066)

* feat(trainer): save checkpoints when load no data * feat(trainer): save checkpoints when load no data * feat(trainer): save checkpoints when load no data * feat(trainer): save checkpoints when load no data
bytedance · Jan 6, 2023 · 86a9f3e · 86a9f3e
1 parent 64faf0a
commit 86a9f3e
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 2 deletions.
diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml
@@ -10,7 +10,12 @@ on:
 
 jobs:
   test:
-    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        # As long as we need Python 3.6 here in the test, we can only use up to Ubuntu 20.
+        # https://github.com/actions/setup-python/issues/544
+        os: [ubuntu-20.04]
+    runs-on: ${{ matrix.os }}
     name: CI tests
     steps:
       - uses: actions/checkout@v2

diff --git a/fedlearner/trainer/trainer_master.py b/fedlearner/trainer/trainer_master.py
@@ -318,7 +318,7 @@ def _session_run(self, estimator):
                 and (self._save_checkpoint_secs \
                     or self._save_checkpoint_steps):
                 hooks.append(
-                    _CheckpointSaverHook(
+                    tf.train.CheckpointSaverHook(
                         checkpoint_dir=self._checkpoint_path,
                         save_secs=self._save_checkpoint_secs,
                         save_steps=self._save_checkpoint_steps,