1. Learn
  2. /
  3. Courses
  4. /
  5. Scalable AI Models with PyTorch Lightning

Connected

Exercise

Introducing the LightningModule

Get ready to build your first LightningModule! In this hands-on exercise, you'll set up the core structure of a classification workflow. You'll define a linear layer, pass data through it in the forward method, and compute the loss in the training step. This clean structure gives you a solid foundation to start experimenting with your models.

The torch and lightning.pytorch, imported as pl, have been preloaded for you.

Instructions

100 XP
  • Define a class LightModel that inherits from pl.LightningModule.
  • Define a linear layer to transform your input, assuming the input features are 16 and there are 10 output classes.