class PyTorchLightningNet(pl.LightningModule):

    def __init__(self):
        # Define model and metrics

    def forward(self, x):
        # Forward pass
        
    def training_step(self, batch, batch_idx):
        # Training step
    
    def validation_step(self, batch, batch_idx):
    
        # Get input and target from the batch
        x, y = batch
        
        # Run forward pass
        y_hat = self.forward(x)
        
        # Compute the loss
        loss = F.cross_entropy(y_hat, y)
        
        # Compute metrics
        accuracy  = self.accuracy(y_hat, y)
        precision = self.precision(y_hat, y)
        
        # Log the loss
        self.log_dict({'test_loss': loss,
                       'val_accuracy':accuracy,
                       'val_precision': precision},
                      on_step=True, on_epoch=True, prog_bar=True)
        
        return loss
pytorch lightning
