How to properly load and save data module

[mfoglio]

I am using checkpointing with the goal of stopping and resuming the training later on. One my main concern, is making sure that the train / val / test split remains the same. I am using a ModelCheckpoint to save the status of the training process. However, after opening a .ckpt file using torch.load I noticed that there is nothing related to the status of my DataModule.
Here is my data module:

import os
from typing import Any, Dict, List, Optional

from torch.utils.data import DataLoader, Dataset
from pytorch_lightning import LightningDataModule
from mycode.dataset.attributes import Attribute


class VehiclesDataModule(LightningDataModule):

    def __init__(self, train_dataset: Dataset, val_dataset: Dataset, test_dataset: Dataset,
                 attributes: List[Attribute], batch_size: int, num_workers: int):
        super().__init__()
        self.train_dataset = train_dataset
        self.val_dataset = val_dataset
        self.test_dataset = test_dataset
        self.attributes = attributes
        self.batch_size = batch_size
        self.num_workers = num_workers

    def setup(self, stage: Optional[str] = None):
        pass

    def _get_data_loaders_kwargs(self):
        return dict(
            batch_size=self.batch_size,
            num_workers=self.num_workers,
            pin_memory=True,
            persistent_workers=True,
            # drop_last=True
            prefetch_factor=2
        )

    def train_dataloader(self):
        return DataLoader(self.train_dataset, **self._get_data_loaders_kwargs())

    def val_dataloader(self):
        return DataLoader(self.val_dataset, **self._get_data_loaders_kwargs())

    def test_dataloader(self):
        return DataLoader(self.test_dataset, **self._get_data_loaders_kwargs())

    def predict_dataloader(self):
        return DataLoader(self.mnist_predict, **self._get_data_loaders_kwargs())

What would be the best way to load and save the data module? Every time I create them the train, val, and test datasets could contain different records. The records are saved as a list of index List[int] as dataset.annotation_ids. Therefore, I need to store this information to resume training. Is there a way to include this in the ModelCheckpoint?

[akihironitta]

@mfoglio Why not simply fixing the seed?

from pytorch_lightning import seed_everything

seed_everything(42)

https://pytorch-lightning.readthedocs.io/en/stable/api/pytorch_lightning.utilities.seed.html?highlight=seed_everything#pytorch_lightning.utilities.seed.seed_everything