Guide · Agent skill
How to use the PyTorch Lightning Training agent skill
Updated · written with AI from the project's own documentation and checked against our catalogue data
The pytorch-lightning agent skill gives your AI agent working knowledge of PyTorch Lightning, a framework that organizes PyTorch training code and removes repetitive boilerplate. It is aimed at people who already write PyTorch models and datasets and want cleaner training code. After this guide you will know what the skill changes, how to install it, and how to ask your agent for Lightning-style training code.
Quick start
- 1Download and unzip the skill. Download the ZIP from this page and unzip it into ~/.claude/skills/ or a project's .claude/skills/ folder. Keep the references folder with SKILL.md.
- 2Install lightning. Make sure Python and PyTorch are installed, then run pip install lightning.
- 3Describe your training task. Ask the agent in plain language to convert or set up your training script, naming Lightning or the Trainer.
- 4Review the generated module. Check that training_step returns the loss, configure_optimizers returns your optimizer, and validation metrics are logged under the names your callbacks monitor.
- 5Test on CPU or one GPU. Run a short job with accelerator='cpu' and devices=1 before enabling multiple devices.
- 6Scale up. Set devices, strategy ('ddp', 'fsdp' or 'deepspeed') or num_nodes once the small run works.
What the skill makes the agent do
Once installed, the agent follows the skill's guidance when it writes or edits your training code. The skill is mainly reference material and code examples. It covers converting plain PyTorch training loops into LightningModule classes, adding validation and testing steps, setting up checkpointing and early stopping, configuring learning rate schedules, and running the same code on one GPU or many GPUs and machines using DDP, FSDP or DeepSpeed.
- Convert a plain PyTorch training loop into a LightningModule
- Add validation and test steps and pass a validation loader to the Trainer
- Set up ModelCheckpoint, EarlyStopping and LearningRateMonitor callbacks
- Configure learning rate schedulers in configure_optimizers
- Switch between single GPU, multi-GPU and multi-node training with Trainer settings
Install the pytorch-lightning skill
Download the ZIP for the skill from this page. To install it for Claude Code, unzip it into ~/.claude/skills/ to make it available everywhere, or into a project's .claude/skills/ folder to limit it to that project. To use it in the Claude apps, upload the ZIP under Settings → Capabilities → Skills.
The skill's own files are SKILL.md plus three reference files: references/callbacks.md, references/distributed.md and references/hyperparameter-tuning.md. Keep them together in the same folder when you unzip.
- Claude Code, all projects: unzip into ~/.claude/skills/
- Claude Code, one project: unzip into the project's .claude/skills/
- Claude apps: Settings → Capabilities → Skills, then upload the ZIP
Install the Python requirements
The skill assumes Python, PyTorch and the lightning Python package. The source gives this install command for Lightning. Optional extras are multiple GPUs for distributed training and TensorBoard for logging.
- Python and PyTorch installed
- Optional: multiple GPUs for distributed training
- Optional: TensorBoard for logging
pip install lightningTrigger the skill with natural requests
You do not need a special command. Describe the training task in plain language and mention Lightning or the Trainer, so the agent knows to use the skill. Give the agent your existing training script when you are converting code, and say which hardware you want to use.
- Convert this PyTorch training loop to a PyTorch Lightning module.
- Add validation and test steps to my Lightning model and log val_loss and val_acc.
- Set up checkpointing that keeps the best 3 models by val_loss and early stopping with patience 5.
- Move this Lightning training script to 8 GPUs using DDP.
- Add a cosine annealing learning rate schedule to my LightningModule.
Example: converting a training loop
The skill's main pattern moves the model, optimizer and step logic into a LightningModule. Device placement is handled by the Trainer, so the training step no longer needs .to('cuda') calls. The optimizer is returned from configure_optimizers, and the Trainer runs the loop. The source reports that this takes a 40-plus line loop down to about 15 lines.
- training_step returns the loss; no manual zero_grad, backward or step calls
- configure_optimizers returns the optimizer
- Trainer(max_epochs=10, accelerator='gpu') runs the loop
trainer = L.Trainer(max_epochs=10, accelerator='gpu')
trainer.fit(LitModel(), train_loader)Example: checkpointing and early stopping
For monitoring, the skill shows ModelCheckpoint and EarlyStopping callbacks passed to the Trainer. The example keeps the best three models by validation loss and stops training if the loss does not improve for five epochs. Review the monitor name and mode, because they must match a metric your model actually logs.
- ModelCheckpoint(monitor='val_loss', mode='min', save_top_k=3)
- EarlyStopping(monitor='val_loss', patience=5, mode='min')
- LearningRateMonitor(logging_interval='epoch') logs the learning rate to TensorBoard
trainer = L.Trainer(max_epochs=100, callbacks=[checkpoint, early_stop, lr_monitor])Test on a small run before scaling up
The skill's troubleshooting advice is to test on CPU first, then move to GPU. Distributed settings such as strategy='ddp' and devices=8 change how processes are launched, so check them on a small run before a long job. Generated distributed configurations can waste expensive GPU time if they are wrong.
- Start with accelerator='cpu', devices=1
- Then try accelerator='gpu', devices=1
- Only then use multiple devices, strategy='ddp' or 'fsdp', or num_nodes
trainer = L.Trainer(accelerator='cpu', devices=1)What people use it for
- Converting an existing research PyTorch script into a LightningModule with a standard structure
- Adding validation, test metrics and best-model checkpointing to an existing training run
- Setting up early stopping and learning rate scheduling without writing the loop yourself
- Moving a working single-GPU training script to multiple GPUs or machines with the same code
Frequently asked questions
How do I install the pytorch-lightning agent skill?
Download the ZIP from this page and unzip it into ~/.claude/skills/ for all projects, or into a project's .claude/skills/ folder for one project. In the Claude apps, upload the ZIP under Settings → Capabilities → Skills.
What do I need to use this skill?
You need Python, PyTorch and the lightning package, installed with pip install lightning. Multiple GPUs are optional and needed only for distributed training, and TensorBoard is optional for logging. The skill assumes you already know PyTorch, datasets and GPUs.
Is the PyTorch Lightning skill free and what is its licence?
The skill is listed under the MIT licence in the source. Check the licence file in the repository for the exact terms before reuse. The skill itself runs no scripts and uses no network access, according to this listing.
Does the skill run code or access the internet?
According to this listing, the skill does not run scripts and does not use the network. It is reference material and code examples the agent follows. The code you run, including any distributed training job, is still yours to launch and check.
Can I use PyTorch Lightning for TensorFlow or JAX projects?
No. The skill covers PyTorch only, so its guidance does not apply to TensorFlow or JAX code. For those, the source suggests alternatives such as Keras for TensorFlow.
How does this compare to Accelerate or Ray Train?
The source says Accelerate suits minimal changes to existing code with more flexibility, and Ray Train suits multi-node orchestration and hyperparameter tuning. Lightning suits teams that want a standard structure, built-in callbacks and logging, and the same code across hardware. Raw PyTorch gives the most control.
New agent skills, repos and AI tools — every week
The best new skills and open-source projects, the fastest-growing repos, and one free prompt. Five minutes to read. Unsubscribe anytime.
Our verdict
Use this skill if you are a machine learning engineer or researcher who already writes PyTorch and wants cleaner, standard training code that scales from one GPU to many. Pick something else if you are new to training models, use TensorFlow or JAX, or need to run distributed jobs unreviewed, since generated multi-GPU configurations should be tested on a small run first.
Download the skill