<jupyter_start><jupyter_text>Three-class classification model<jupyter_code>class SimpleClassifier(pl.LightningModule):

    def __init__(self):
        super().__init__()

        # Linear  
        self.layer_1 = torch.nn.Linear(4, 32)
        self.layer_2 = torch.nn.Linear(32, 64)
        self.layer_3 = torch.nn.Linear(64, 32)
        self.layer_4 = torch.nn.Linear(32, 2)

    def forward(self, x):

        # Layer 1 
        x = self.layer_1(x)
        x = torch.relu(x)

        # Layer 2
        x = self.layer_2(x)
        x = torch.relu(x)
        
        # Layer 3
        x = self.layer_3(x)
        x = torch.relu(x)
        
        # Layer 4
        x = self.layer_4(x)
        x = torch.relu(x)

        # Probability distribution over labels
        x = torch.log_softmax(x, dim=1)

        return x

    def cross_entropy_loss(self, logits, labels):
        return F.nll_loss(logits, labels)

    def training_step(self, train_batch, batch_idx):
        x, y = train_batch
        logits = self.forward(x)
        loss = self.cross_entropy_loss(logits, y)
        self.log('train_loss', loss)
        return loss

    def validation_step(self, val_batch, batch_idx):
        x, y = val_batch
        logits = self.forward(x)
        loss = self.cross_entropy_loss(logits, y)
        self.log('val_loss', loss)

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=1e-2)
        return optimizer
class MyDataModule(pl.LightningDataModule):
    
    def __init__(self, X, y, train_size):
        super().__init__()
        self.X = X
        self.y = y
        self.train_size = train_size

    def setup(self, stage):
        
        X_train = torch.tensor(self.X[:self.train_size])
        y_train = torch.tensor(self.y[:self.train_size], dtype=torch.long)
        X_test = torch.tensor(self.X[self.train_size:])
        y_test = torch.tensor(self.y[self.train_size:], dtype=torch.long)
        
        self.train = TensorDataset(X_train, y_train)
        self.test = TensorDataset(X_test, y_test)

    def train_dataloader(self):
        return DataLoader(self.train, batch_size=16)

    def val_dataloader(self):
        return DataLoader(self.test, batch_size=16)
X = np.array([
    [1, 1, 1, 1],
    [2, 4, 6, 24], 
    [-1, -2, -1, -192],
    [-191, -3, -2, -7],
    [102, 12, 16, 200],
    [7, 9, 13, 177],
] * 200, dtype='float32')

y = np.array([
    0, 
    0, 
    1, 
    1, 
    0, 
    0
] * 200, dtype='int32')
 

data_module = MyDataModule(X, y, train_size=16)
early_stopping = EarlyStopping('val_loss', patience=20)

model = SimpleClassifier()

trainer = pl.Trainer(max_epochs=1000, callbacks=[early_stopping])

trainer.fit(model, data_module)
model(torch.tensor([[1, 1, 1, 79]], dtype=torch.float)).argmax()<jupyter_output><empty_output>