    def retrieve_and_delete_saved(self, shard: Optional[int] = None):
        """
        Helper function for the tests below. Finds the weight and training state files in
        self.TEST_DIR, parses their names for the epochs that were saved, deletes them,
        and returns the saved epochs as two lists of integers.
        """
        serialization_files = os.listdir(self.TEST_DIR)

        model_checkpoints = [x for x in serialization_files if "model_state_" in x]
        if shard is not None:
            model_checkpoints = [x for x in model_checkpoints if x.endswith(f"_w{shard}.th")]
        found_model_states = [Checkpointer._parse_model_state_path(x) for x in model_checkpoints]
        for f in model_checkpoints:
            os.remove(os.path.join(self.TEST_DIR, f))

        training_checkpoints = [x for x in serialization_files if "training_state_" in x]
        if shard is not None:
            training_checkpoints = [x for x in training_checkpoints if x.endswith(f"_w{shard}.th")]
        found_training_states = [
            Checkpointer._parse_training_state_path(x) for x in training_checkpoints
        ]
        for f in training_checkpoints:
            os.remove(os.path.join(self.TEST_DIR, f))
        return sorted(found_model_states), sorted(found_training_states)
Helper function for the tests below. Finds the weight and training state files in
self.TEST_DIR, parses their names for the epochs that were saved, deletes them,
and returns the saved epochs as two lists of integers.
tests/training/checkpointer_test.py