Validation step is ignored when using DataModule
What is your question?
Hi, guys!
I created my own DataModule and loading it to the trainer. However, it appears that the "fit" is skipping the validation step.
How can I ensure that the code runs through the validation step too?
Code
class DataModule(pl.LightningDataModule):
    def __init__(self, batch_size=25, seed=0):
    # def __init__(self, dataset, batch_size=25, seed=0):
        super().__init__()
        self.dataset = dataset
        self.batch_size = batch_size
        self.seed = seed
        self.split = [801, 100, 100]
        # self.transform = torchvision.transforms.ToTensor()

    def setup(self, stage=None):
        # train/valid/test split
        # and assign to use in dataloaders via self
        train_set, valid_set, test_set = torch.utils.data.random_split(self.dataset, self.split, generator=torch.Generator().manual_seed(self.seed))

        if stage == 'fit' or stage is None:

            self.train_set = train_set
            self.valid_set = valid_set

        if stage == 'test' or stage is None:
            self.test_set = test_set

    def train_dataloader(self):
        return torch.utils.data.DataLoader(self.train_set, batch_size=self.batch_size, shuffle=True)

    def valid_dataloader(self):
        return torch.utils.data.DataLoader(self.valid_set, batch_size=self.batch_size, shuffle=False)

    def test_dataloader(self):
        return torch.utils.data.DataLoader(self.test_set, batch_size=self.batch_size, shuffle=False)

class LitReg(pl.LightningModule):
    def __init__(self, in_dims, out_dims, lr=2e-4, max_dict={}):
        super().__init__()
        self.in_size = in_dims
        self.out_size = out_dims
        self.lr = lr
        self.max_dict = max_dict

        # model
        self.model = nn.Sequential(
            nn.Linear(self.in_size, self.in_size),

            nn.LeakyReLU(0.02),

            nn.Linear(self.in_size, self.out_size)
        )

        self.model.apply(self.weights_init)

    def forward(self, data):
        out = self.model(data)
        return out

    def training_step(self, batch, batch_idx):
        x, y, l_rate = batch
        pred_y = self.model(x)
        train_loss = F.mse_loss(pred_y, y) 

        self.log('train_loss', train_loss, prog_bar=True)
        return train_loss

    def validation_step(self, batch, batch_idx):
        self._shared_eval(batch, batch_idx, 'val')

    def test_step(self, batch, batch_idx):
        self._shared_eval(batch, batch_idx, 'test')

    def _shared_eval(self, batch, batch_idx, prefix):
        x, y, l_rate = batch
        pred_y = self.model(x)

        loss = F.mse_loss(pred_y, y) 
        self.log(f'{prefix}_loss', loss, prog_bar=True)
        return loss

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), self.lr)
        return optimizer

    def weights_init(self, m):
        if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d):
            torch.nn.init.normal_(m.weight, 0.0, 0.02)
        if isinstance(m, nn.BatchNorm2d) or isinstance(m, nn.BatchNorm1d):
            torch.nn.init.normal_(m.weight, 0.0, 0.02)
            torch.nn.init.constant_(m.bias, 0)
        if isinstance(m, nn.Linear):
            torch.nn.init.normal_(m.weight, 0.0, 0.02)
            torch.nn.init.constant_(m.bias, 0)

What have you tried?
Placing breakpoints to debug in VSCode, but all in vain.
Also accessed both valid and test datasets and loaders. All looks set.
What is working? If I load the data the following way.
    **train_loader = DataLoader(X_train, batch_size=args.batch_size)
    val_loader = DataLoader(X_val, batch_size=args.batch_size)
    test_loader = DataLoader(X_test, batch_size=args.batch_size)

    trainer.fit(model, train_loader, val_loader)**

What's your environment?

OS: Win
Packaging pip
Version 1.1.3

Thank you for your attention!