---
title: API Reference
description: API reference for all public classes and functions.
canonical: 'https://lightning.ai/docs/pytorch/stable/optional-api/api_references'
---
# API Reference {#api-reference}

## accelerators {#accelerators}

- [`Accelerator`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.accelerators.Accelerator.md) — The Accelerator base class for Lightning PyTorch.
- [`CPUAccelerator`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.accelerators.CPUAccelerator.md) — Accelerator for CPU devices.
- [`CUDAAccelerator`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.accelerators.CUDAAccelerator.md) — Accelerator for NVIDIA CUDA devices.
- [`XLAAccelerator`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.accelerators.XLAAccelerator.md) — Accelerator for XLA devices, normally TPUs.

## callbacks {#callbacks}

- [`BackboneFinetuning`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.callbacks.BackboneFinetuning.md) — Finetune a backbone model based on a learning rate user-defined scheduling.
- [`BaseFinetuning`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.callbacks.BaseFinetuning.md) — This class implements the base logic for writing your own Finetuning Callback.
- [`BasePredictionWriter`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.callbacks.BasePredictionWriter.md) — Base class to implement how the predictions should be stored.
- [`BatchSizeFinder`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.callbacks.BatchSizeFinder.md) — Finds the largest batch size supported by a given model before encountering an out of memory (OOM) error.
- [`Callback`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.callbacks.Callback.md) — Abstract base class used to build new callbacks.
- [`DeviceStatsMonitor`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.callbacks.DeviceStatsMonitor.md) — Automatically monitors and logs device stats during training, validation and testing stage.
DeviceStatsMonitor is a special callback as it requires a logger to passed as argument to the Trainer.
- [`EarlyStopping`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.callbacks.EarlyStopping.md) — Monitor a metric and stop training when it stops improving.
- [`GradientAccumulationScheduler`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.callbacks.GradientAccumulationScheduler.md) — Change gradient accumulation factor according to scheduling.
- [`LambdaCallback`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.callbacks.LambdaCallback.md) — Create a simple callback on the fly using lambda functions.
- [`LearningRateFinder`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.callbacks.LearningRateFinder.md) — The LearningRateFinder callback enables the user to do a range test of good initial learning rates, to
reduce the amount of guesswork in picking a good starting learning rate.
- [`LearningRateMonitor`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.callbacks.LearningRateMonitor.md) — Automatically monitor and logs learning rate for learning rate schedulers during training.
- [`ModelCheckpoint`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.callbacks.ModelCheckpoint.md) — Save the model after every epoch by monitoring a quantity. Every logged metrics are passed to the
Logger for the version it gets saved in the same directory as the
checkpoint.
- [`ModelPruning`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.callbacks.ModelPruning.md) — Model pruning Callback, using PyTorch’s prune utilities. This callback is responsible of pruning networks
parameters during training.
- [`ModelSummary`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.callbacks.ModelSummary.md) — Generates a summary of all layers in a LightningModule.
- [`OnExceptionCheckpoint`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.callbacks.OnExceptionCheckpoint.md) — Used to save a checkpoint on exception.
- [`ProgressBar`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.callbacks.ProgressBar.md) — The base class for progress bars in Lightning. It is a Callback that keeps
track of the batch progress in the Trainer. You should implement your
highly custom progress bars with this as the base class
- [`RichModelSummary`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.callbacks.RichModelSummary.md) — Generates a summary of all layers in a LightningModule with rich text
formatting.
- [`RichProgressBar`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.callbacks.RichProgressBar.md) — Create a progress bar with rich text formatting.
- [`StochasticWeightAveraging`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.callbacks.StochasticWeightAveraging.md) — Implements the Stochastic Weight Averaging (SWA) Callback to average a model.
- [`SpikeDetection`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.callbacks.SpikeDetection.md)
- [`ThroughputMonitor`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.callbacks.ThroughputMonitor.md) — Computes and logs throughput with the Throughput
- [`Timer`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.callbacks.Timer.md) — The Timer callback tracks the time spent in the training, validation, and test loops and interrupts the Trainer
if the given time limit for the training loop is reached.
- [`TQDMProgressBar`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.callbacks.TQDMProgressBar.md) — This is the default progress bar used by Lightning. It prints to stdout using the tqdm package and
shows up to four different bars:
- [`WeightAveraging`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.callbacks.WeightAveraging.md) — A callback that updates an averaged model for Stochastic Weight Averaging (SWA) or Exponential Moving Average
(EMA) after each training step.

## cli {#cli}

- [`LightningCLI`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.cli.LightningCLI.md) — Implementation of a configurable command line tool for pytorch-lightning.
- [`LightningArgumentParser`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.cli.LightningArgumentParser.md) — Extension of jsonargparse’s ArgumentParser for pytorch-lightning.
- [`SaveConfigCallback`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.cli.SaveConfigCallback.md) — Saves a LightningCLI config to the log_dir when training starts.

## core {#core}

- [`CheckpointHooks`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.core.hooks.CheckpointHooks.md) — Hooks to be used with Checkpointing.
- [`DataHooks`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.core.hooks.DataHooks.md) — Hooks to be used for data related stuff.
- [`ModelHooks`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.core.hooks.ModelHooks.md) — Hooks to be used in LightningModule.
- [`LightningDataModule`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.core.LightningDataModule.md) — A DataModule standardizes the training, val, test splits, data preparation and transforms. The main advantage is
consistent data splits, data preparation and transforms across models.
- [`LightningModule`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.core.LightningModule.md)
- [`HyperparametersMixin`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.core.mixins.HyperparametersMixin.md)
- [`LightningOptimizer`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.core.optimizer.LightningOptimizer.md) — This class is used to wrap the user optimizers and handle properly the backward and optimizer_step logic across
accelerators, AMP, accumulate_grad_batches.

## loggers {#loggers}

- [`logger`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.loggers.logger.md) — Dummy logger for internal use.
- [`litlogger`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.loggers.litlogger.md) — Logger that enables remote experiment tracking, logging, and artifact management on lightning.ai.
- [`comet`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.loggers.comet.md) — Track your parameters, metrics, source code and more using Comet.
- [`csv_logs`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.loggers.csv_logs.md) — Log to local file system in yaml and CSV format.
- [`mlflow`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.loggers.mlflow.md) — Log using MLflow.
- [`tensorboard`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.loggers.tensorboard.md) — Log to local or remote file system in TensorBoard format.
- [`wandb`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.loggers.wandb.md) — Log using Weights and Biases.

### plugins {#plugins}

#### precision {#precision}

- [`DeepSpeedPrecision`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.plugins.precision.DeepSpeedPrecision.md) — Precision plugin for DeepSpeed integration.
- [`DoublePrecision`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.plugins.precision.DoublePrecision.md) — Plugin for training with double (torch.float64) precision.
- [`HalfPrecision`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.plugins.precision.HalfPrecision.md) — Plugin for training with half precision.
- [`FSDPPrecision`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.plugins.precision.FSDPPrecision.md) — Precision plugin for training with Fully Sharded Data Parallel (FSDP).
- [`MixedPrecision`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.plugins.precision.MixedPrecision.md) — Plugin for Automatic Mixed Precision (AMP) training with torch.autocast.
- [`Precision`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.plugins.precision.Precision.md) — Base class for all plugins handling the precision-specific parts of the training.
- [`XLAPrecision`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.plugins.precision.XLAPrecision.md) — Plugin for training with XLA.
- [`TransformerEnginePrecision`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.plugins.precision.TransformerEnginePrecision.md) — Plugin for training with fp8 precision via nvidia’s
Transformer Engine.
- [`BitsandbytesPrecision`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.plugins.precision.BitsandbytesPrecision.md) — Plugin for quantizing weights with bitsandbytes.

#### environments {#environments}

- [`ClusterEnvironment`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.plugins.environments.ClusterEnvironment.md) — Specification of a cluster environment.
- [`KubeflowEnvironment`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.plugins.environments.KubeflowEnvironment.md) — Environment for distributed training using the PyTorchJob operator from Kubeflow.
- [`LightningEnvironment`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.plugins.environments.LightningEnvironment.md) — The default environment used by Lightning for a single node or free cluster (not managed).
- [`LSFEnvironment`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.plugins.environments.LSFEnvironment.md) — An environment for running on clusters managed by the LSF resource manager.
- [`MPIEnvironment`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.plugins.environments.MPIEnvironment.md) — An environment for running on clusters with processes created through MPI.
- [`SLURMEnvironment`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.plugins.environments.SLURMEnvironment.md) — Cluster environment for training on a cluster managed by SLURM.
- [`TorchElasticEnvironment`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.plugins.environments.TorchElasticEnvironment.md) — Environment for fault-tolerant and elastic training with torchelastic
- [`XLAEnvironment`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.plugins.environments.XLAEnvironment.md) — Cluster environment for training on a TPU Pod with the PyTorch/XLA library.

#### io {#io}

- [`AsyncCheckpointIO`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.plugins.io.AsyncCheckpointIO.md) — AsyncCheckpointIO enables saving the checkpoints asynchronously in a thread.
- [`CheckpointIO`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.plugins.io.CheckpointIO.md) — Interface to save/load checkpoints as they are saved through the Strategy.
- [`TorchCheckpointIO`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.plugins.io.TorchCheckpointIO.md) — CheckpointIO that utilizes torch.save() and torch.load() to save and load checkpoints respectively,
common for most use cases.
- [`XLACheckpointIO`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.plugins.io.XLACheckpointIO.md) — CheckpointIO that utilizes xm.save to save checkpoints for TPU training strategies.

#### others {#others}

- [`LayerSync`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.plugins.LayerSync.md) — Abstract base class for creating plugins that wrap layers of a model with synchronization logic for
multiprocessing.
- [`TorchSyncBatchNorm`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.plugins.TorchSyncBatchNorm.md) — A plugin that wraps all batch normalization layers of a model with synchronization logic for multiprocessing.

## profiler {#profiler}

- [`AdvancedProfiler`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.profilers.AdvancedProfiler.md) — This profiler uses Python’s cProfiler to record more detailed information about time spent in each function call
recorded during a given action.
- [`PassThroughProfiler`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.profilers.PassThroughProfiler.md) — This class should be used when you don’t want the (small) overhead of profiling.
- [`Profiler`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.profilers.Profiler.md) — If you wish to write a custom profiler, you should inherit from this class.
- [`PyTorchProfiler`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.profilers.PyTorchProfiler.md) — This profiler uses PyTorch’s Autograd Profiler and lets you inspect the cost of
different operators inside your model - both on the CPU and GPU.
- [`SimpleProfiler`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.profilers.SimpleProfiler.md) — This profiler simply records the duration of actions (in seconds) and reports the mean duration of each action
and the total time spent over the entire training run.
- [`XLAProfiler`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.profilers.XLAProfiler.md) — XLA Profiler will help you debug and optimize training workload performance for your models using Cloud TPU
performance tools.

## trainer {#trainer}

- [`Trainer`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.trainer.trainer.Trainer.md) — Customize every aspect of training via flags.

## strategies {#strategies}

- [`DDPStrategy`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.strategies.DDPStrategy.md) — Strategy for multi-process single-device training on one or multiple nodes.
- [`DeepSpeedStrategy`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.strategies.DeepSpeedStrategy.md) — Provides capabilities to run training using the DeepSpeed library, with training optimizations for large
billion parameter models. For more information: DeepSpeed.
- [`FSDPStrategy`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.strategies.FSDPStrategy.md) — Strategy for Fully Sharded Data Parallel provided by torch.distributed.
- [`ModelParallelStrategy`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.strategies.ModelParallelStrategy.md) — Enables user-defined parallelism applied to a model.
- [`ParallelStrategy`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.strategies.ParallelStrategy.md) — Strategy for training with multiple processes in parallel.
- [`SingleDeviceStrategy`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.strategies.SingleDeviceStrategy.md) — Strategy that handles communication on a single device.
- [`SingleDeviceXLAStrategy`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.strategies.SingleDeviceXLAStrategy.md) — Strategy for training on a single XLA device.
- [`Strategy`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.strategies.Strategy.md) — Base class for all strategies that change the behaviour of the training, validation and test- loop.
- [`XLAStrategy`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.strategies.XLAStrategy.md) — Strategy for training multiple TPU devices using the torch_xla.distributed.xla_multiprocessing.spawn()
method.

## tuner {#tuner}

- [`Tuner`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.tuner.tuning.Tuner.md) — Tuner class to tune your model.

## utilities {#utilities}

- [`combined_loader`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.utilities.combined_loader.md) — Combines different iterables under specific sampling modes.
- [`data`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.utilities.data.md) — Unpack a batch to find a torch.Tensor.
- [`deepspeed`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.utilities.deepspeed.md) — Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated state_dict file that can be loaded with
torch.load(file) + load_state_dict() and used for training without DeepSpeed. It gets copied into
- [`memory`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.utilities.memory.md) — Garbage collection Torch (CUDA) memory.
- [`model_summary`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.utilities.model_summary.md)
- [`parsing`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.utilities.parsing.md) — Extended dictionary accessible with dot notation.
- [`rank_zero`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.utilities.rank_zero.md)
- [`seed`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.utilities.seed.md) — A context manager that resets the global random state on exit to what it was before entering.
- [`warnings`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.utilities.warnings.md)

### `lightning.pytorch.utilities.measure_flops(model, forward_fn, loss_fn=None)` {#lightning.pytorch.utilities.measure_flops}

Utility to compute the total number of FLOPs used by a module during training or during inference.

It’s recommended to create a meta-device model for this:

Example:

```python
with torch.device("meta"):
    model = MyModel()
    x = torch.randn(2, 32)

model_fwd = lambda: model(x)
fwd_flops = measure_flops(model, model_fwd)

model_loss = lambda y: y.sum()
fwd_and_bwd_flops = measure_flops(model, model_fwd, model_loss)
```

Parameters:

*   **model** ([`Module`](https://docs.pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module "(in PyTorch v2.15)")) – The model whose FLOPs should be measured.
    
*   **forward\_fn** ([`Callable`](https://docs.python.org/3/library/typing.html#typing.Callable "(in Python v3.15)")\[\[\], [`Tensor`](https://docs.pytorch.org/docs/stable/tensors.html#torch.Tensor "(in PyTorch v2.15)")\]) – A function that runs `forward` on the model and returns the result.
    
*   **loss\_fn** ([`Optional`](https://docs.python.org/3/library/typing.html#typing.Optional "(in Python v3.15)")\[[`Callable`](https://docs.python.org/3/library/typing.html#typing.Callable "(in Python v3.15)")\[\[[`Tensor`](https://docs.pytorch.org/docs/stable/tensors.html#torch.Tensor "(in PyTorch v2.15)")\], [`Tensor`](https://docs.pytorch.org/docs/stable/tensors.html#torch.Tensor "(in PyTorch v2.15)")\]\]) – A function that computes the loss given the `forward_fn` output. If provided, the loss and backward FLOPs will be included in the result.
    

Return type:

[`int`](https://docs.python.org/3/builtins/functions.html#int "(in Python v3.15)")

---

**Navigation**

- Parent: [Optional API](/docs/pytorch/2.6.6/optional-api.md)