PyTorch Lightning Training Loop

import torch
import torch.nn as nn
import pytorch_lightning as pl
from torch.utils.data import DataLoader, TensorDataset

class SimpleClassifier(pl.LightningModule):
    def __init__(self, input_dim: int = 20, hidden: int = 64, num_classes: int = 3):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(input_dim, hidden), nn.ReLU(), nn.Dropout(0.2), nn.Linear(hidden, num_classes))
        self.criterion = nn.CrossEntropyLoss()

    def forward(self, x):
        return self.net(x)

    def training_step(self, batch, batch_idx):
        x, y = batch
        loss = self.criterion(self(x), y)
        self.log('train_loss', loss)
        return loss

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=1e-3)

X = torch.randn(300, 20)
y = torch.randint(0, 3, (300,))
loader  = DataLoader(TensorDataset(X, y), batch_size=32, shuffle=True)
trainer = pl.Trainer(max_epochs=5, enable_progress_bar=True)
trainer.fit(SimpleClassifier(), loader)