How to use the PyTorch Lightning Training agent skill

What the skill makes the agent do

Agent skills are folders of instructions and helper files that an agent loads when a request matches them. This one covers the PyTorch Lightning framework (the lightning package). It helps the agent split training code into a LightningModule for the model and its training, validation, test and prediction steps, a LightningDataModule for data loading, and a Trainer that runs the loop, device management, checkpoints and logging.

The skill is guidance and boilerplate for writing and running training code. It is not a trained model. Its templates use synthetic random data and are only smoke tests, so they do not show scientific accuracy.

• Organises a PyTorch project into LightningModule, LightningDataModule and Trainer
• Starts from bundled template files: scripts/template_lightning_module.py, scripts/template_datamodule.py and scripts/quick_trainer_setup.py
• Points to reference notes on callbacks, logging, the Trainer and distributed training (DDP, FSDP, DeepSpeed)

How to install the PyTorch Lightning skill

You can install the skill in one of two ways. Download the ZIP from this page and unzip it into your personal skills folder, ~/.claude/skills/, or into a project's .claude/skills/ folder if you want it to apply only to that project. Alternatively, upload the ZIP in the Claude apps under Settings → Capabilities → Skills.

The folder you unzip should contain SKILL.md, the references/ folder and the scripts/ folder, so the agent can read the notes and copy the templates.

• Option A: unzip into ~/.claude/skills/ for all your projects
• Option B: unzip into a project's .claude/skills/ folder for that project only
• Option C: upload the ZIP in the Claude apps under Settings → Capabilities → Skills

Install the Lightning package it depends on

The skill expects Python 3.10 or newer, the lightning 2.6.6 package and a compatible PyTorch install. The skill's own installation instructions use uv. Lightning exposes the namespace import lightning as L; the separate pytorch-lightning distribution uses import pytorch_lightning as pl. Do not mix the two namespaces in one application.

Optional extras and logger packages are installed separately, only if you use them.

• Requires Python 3.10 or newer
• Requires lightning 2.6.6 and a compatible PyTorch
• Optional: wandb or mlflow, only if you use those loggers

How to trigger the skill with natural requests

You do not need to name the skill. Describe the task in plain language that matches what the skill covers. The skill is meant to be used when you are building or training a network with PyTorch Lightning, organising code into LightningModules, configuring Trainers for multi-GPU or TPU, building data pipelines with LightningDataModules, or adding callbacks, logging and distributed strategies.

Be specific about your model, data and hardware. The more the agent knows about your setup, the less it has to guess.

• "Convert my PyTorch training loop into a LightningModule with a Trainer."
• "Wrap my image folder loading in a LightningDataModule with train and validation loaders."
• "Set up a Trainer that checkpoints the best model on val/loss and stops early."
• "I have 4 GPUs. Which Lightning strategy should I use for my model, and how do I configure it?"

Examples of a smoke test you can run

The skill's quick workflow copies the three template scripts into a working directory and runs a CPU smoke test. The example below comes from the skill's own instructions. It uses synthetic data, so the result shows that the code runs, not that the model is good. The skill also notes that a fresh model and Trainer should be used when resuming from a checkpoint, and that random labels test mechanics only.

Helpers monitor val/loss by default. Check the bundled scripts before you change the code, because the templates are meant as starting points.

• Run on CPU with accelerator="cpu" and precision="32-true"
• Logs go to logs/ and checkpoints to checkpoints/ in the working directory

What the skill needs and what it does not do

The skill runs scripts, so the agent can execute Python code on your machine. It does not use the network. CPU examples need no network after installation. GPU, DDP, FSDP and DeepSpeed setups need suitable hardware, and the skill notes that these configurations are illustrative and not validated by executed training.

Online loggers need extra packages and credentials, and they send data to external services. The bundled helpers log to local CSV files. NeptuneLogger was removed in Lightning 2.6.4, so do not rely on it.

• Runs scripts: yes. Uses network: no
• Optional logger packages (wandb, mlflow) and credentials are needed only for online logging
• Multi-GPU, FSDP and DeepSpeed need suitable hardware; check before relying on them
• Cloud logging or artifact uploads should be a deliberate decision, since credentials and network access would expose data

Tips for getting good results

Start with a fast check before a real run. The skill recommends fast_dev_run=True on the Trainer to test with one batch. Set a seed with L.seed_everything(seed, workers=True) so runs are reproducible, and use deterministic=True on the Trainer where you need it.

For distributed training, choose the strategy from measured memory and throughput rather than a fixed size rule. Use DDP when each device can hold the full model, optimizer state and activations. Consider FSDP or DeepSpeed only when memory pressure is measured and those features fit your CUDA stack.

When you log a scalar metric across several GPUs, pass sync_dist=True to self.log, because the default is False. For metrics that do not add up across batches, such as AUROC, use a TorchMetrics object rather than averaging per-GPU values.

• Use self.device instead of calling .cuda() so code stays device agnostic
• Call self.save_hyperparameters() in __init__() to save hyperparameters with checkpoints
• Split your data by the real unit of independence (patient, site, batch) before preprocessing, and keep a test set for the final comparison