Fix set_epoch not getting called for prediction dataloaders (#16785)

src/pytorch_lightning/CHANGELOG.md

@@ -5,6 +5,13 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 
+## [1.9.3] - YYYY-MM-DD
+
+### Fixed
+
+- Fixed bug where `set_epoch` was not called for prediction dataloaders ([#16785](https://github.com/Lightning-AI/lightning/pull/16785))
+
+
 ## [1.9.2] - 2023-02-15
 
 ### Changed

src/pytorch_lightning/overrides/distributed.py

@@ -147,3 +147,7 @@ def batch_size(self) -> int:
     @property
     def sampler(self) -> Union[Sampler, Iterable]:
         return self._sampler.sampler
+
+    def set_epoch(self, epoch: int) -> None:
+        if hasattr(self._sampler, "set_epoch"):
+            self._sampler.set_epoch(epoch)

tests/tests_pytorch/loops/test_prediction_loop.py

@@ -0,0 +1,35 @@
+# Copyright The Lightning AI team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+from unittest import mock
+from unittest.mock import call
+
+from pytorch_lightning import Trainer
+from pytorch_lightning.demos.boring_classes import BoringModel
+
+
+def test_prediction_loop_batch_sampler_set_epoch_called(tmp_path):
+    """Tests that set_epoch is called on the dataloader's batch sampler (if any) during prediction."""
+    model = BoringModel()
+    trainer = Trainer(
+        default_root_dir=tmp_path,
+        limit_predict_batches=1,
+        enable_model_summary=False,
+        enable_checkpointing=False,
+        logger=False,
+    )
+    trainer.fit_loop.epoch_progress.current.processed = 2
+
+    with mock.patch("pytorch_lightning.overrides.distributed.IndexBatchSamplerWrapper.set_epoch") as set_epoch_mock:
+        trainer.predict(model)
+    assert set_epoch_mock.mock_calls == [call(2)]