Fixes

Author: carmocca

tests/loops/test_loop_state_dict.py

@@ -11,6 +11,8 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import os
+from unittest import mock
 from unittest.mock import Mock
 
 import pytest
@@ -37,6 +39,7 @@ def test_loops_state_dict():
     assert fit_loop.state_dict() == new_fit_loop.state_dict()
 
 
+@mock.patch.dict(os.environ, {"PL_FAULT_TOLERANT_TRAINING": "1"})
 def test_loops_state_dict_structure():
     trainer = Trainer()
     trainer.train_dataloader = Mock()

tests/models/test_restore.py

@@ -251,7 +251,7 @@ def on_pretrain_routine_end(self) -> None:
     trainer.fit(TestModel(), ckpt_path=ckpt)
     assert trainer.current_epoch == max_epochs
     # TODO(@carmocca): should not need `+1`
-    # assert trainer.global_step == max_epochs * train_batches + 1
+    assert trainer.global_step == max_epochs * train_batches + 1
 
 
 def test_fit_twice(tmpdir):