LightningModule — PyTorch Lightning 2.5.1.post0 documentation (original) (raw)

A LightningModule organizes your PyTorch code into 6 sections:

Initialization (__init__ and setup()).
Train Loop (training_step())
Validation Loop (validation_step())
Test Loop (test_step())
Prediction Loop (predict_step())
Optimizers and LR Schedulers (configure_optimizers())

When you convert to use Lightning, the code IS NOT abstracted - just organized. All the other code that’s not in the LightningModulehas been automated for you by the Trainer.

net = MyLightningModuleNet() trainer = Trainer() trainer.fit(net)

There are no .cuda() or .to(device) calls required. Lightning does these for you.

don't do in Lightning
x = torch.Tensor(2, 3) x = x.cuda() x = x.to(device)

do this instead
x = x # leave it alone!

or to init a new tensor
new_x = torch.Tensor(2, 3) new_x = new_x.to(x)

When running under a distributed strategy, Lightning handles the distributed sampler for you by default.

Don't do in Lightning...
data = MNIST(...) sampler = DistributedSampler(data) DataLoader(data, sampler=sampler)

do this instead
data = MNIST(...) DataLoader(data)

A LightningModule is a torch.nn.Module but with added functionality. Use it as such!

net = Net.load_from_checkpoint(PATH) net.freeze() out = net(x)

Thus, to use Lightning, you just need to organize your code which takes about 30 minutes, (and let’s be real, you probably should do anyway).

Starter Example¶

Here are the only required methods.

import lightning as L import torch

from lightning.pytorch.demos import Transformer

class LightningTransformer(L.LightningModule): def init(self, vocab_size): super().init() self.model = Transformer(vocab_size=vocab_size)

def forward(self, inputs, target):
    return self.model(inputs, target)

def training_step(self, batch, batch_idx):
    inputs, target = batch
    output = self(inputs, target)
    loss = torch.nn.functional.nll_loss(output, target.view(-1))
    return loss

def configure_optimizers(self):
    return torch.optim.SGD(self.model.parameters(), lr=0.1)

Which you can train by doing:

from lightning.pytorch.demos import WikiText2 from torch.utils.data import DataLoader

dataset = WikiText2() dataloader = DataLoader(dataset) model = LightningTransformer(vocab_size=dataset.vocab_size)

trainer = L.Trainer(fast_dev_run=100) trainer.fit(model=model, train_dataloaders=dataloader)

The LightningModule has many convenient methods, but the core ones you need to know about are:

Name	Description
__init__ and setup()	Define initialization here
forward()	To run data through your model only (separate from training_step)
training_step()	the complete training step
validation_step()	the complete validation step
test_step()	the complete test step
predict_step()	the complete prediction step
configure_optimizers()	define optimizers and LR schedulers

Training¶

Training Loop¶

To activate the training loop, override the training_step() method.

class LightningTransformer(L.LightningModule): def init(self, vocab_size): super().init() self.model = Transformer(vocab_size=vocab_size)

def training_step(self, batch, batch_idx):
    inputs, target = batch
    output = self.model(inputs, target)
    loss = torch.nn.functional.nll_loss(output, target.view(-1))
    return loss

Under the hood, Lightning does the following (pseudocode):

enable gradient calculation

torch.set_grad_enabled(True)

for batch_idx, batch in enumerate(train_dataloader): loss = training_step(batch, batch_idx)

# clear gradients
optimizer.zero_grad()

# backward
loss.backward()

# update parameters
optimizer.step()

Train Epoch-level Metrics¶

If you want to calculate epoch-level metrics and log them, use log().

def training_step(self, batch, batch_idx): inputs, target = batch output = self.model(inputs, target) loss = torch.nn.functional.nll_loss(output, target.view(-1))

# logs metrics for each training_step,
# and the average across the epoch, to the progress bar and logger
self.log("train_loss", loss, on_step=True, on_epoch=True, prog_bar=True, logger=True)
return loss

The log() method automatically reduces the requested metrics across a complete epoch and devices. Here’s the pseudocode of what it does under the hood:

outs = [] for batch_idx, batch in enumerate(train_dataloader): # forward loss = training_step(batch, batch_idx) outs.append(loss.detach())

# clear gradients
optimizer.zero_grad()
# backward
loss.backward()
# update parameters
optimizer.step()

note: in reality, we do this incrementally, instead of keeping all outputs in memory

epoch_metric = torch.mean(torch.stack(outs))

Train Epoch-level Operations¶

In the case that you need to make use of all the outputs from each training_step(), override the on_train_epoch_end() method.

class LightningTransformer(L.LightningModule): def init(self, vocab_size): super().init() self.model = Transformer(vocab_size=vocab_size) self.training_step_outputs = []

def training_step(self, batch, batch_idx):
    inputs, target = batch
    output = self.model(inputs, target)
    loss = torch.nn.functional.nll_loss(output, target.view(-1))
    preds = ...
    self.training_step_outputs.append(preds)
    return loss

def on_train_epoch_end(self):
    all_preds = torch.stack(self.training_step_outputs)
    # do something with all preds
    ...
    self.training_step_outputs.clear()  # free memory

Validation¶

Validation Loop¶

To activate the validation loop while training, override the validation_step() method.

class LightningTransformer(L.LightningModule): def validation_step(self, batch, batch_idx): inputs, target = batch output = self.model(inputs, target) loss = F.cross_entropy(y_hat, y) self.log("val_loss", loss)

Under the hood, Lightning does the following (pseudocode):

...

for batch_idx, batch in enumerate(train_dataloader): loss = model.training_step(batch, batch_idx) loss.backward() # ...

if validate_at_some_point:
    # disable grads + batchnorm + dropout
    torch.set_grad_enabled(False)
    model.eval()

    # ----------------- VAL LOOP ---------------
    for val_batch_idx, val_batch in enumerate(val_dataloader):
        val_out = model.validation_step(val_batch, val_batch_idx)
    # ----------------- VAL LOOP ---------------

    # enable grads + batchnorm + dropout
    torch.set_grad_enabled(True)
    model.train()

You can also run just the validation loop on your validation dataloaders by overriding validation_step()and calling validate().

model = LightningTransformer(vocab_size=dataset.vocab_size) trainer = L.Trainer() trainer.validate(model)

Note

It is recommended to validate on single device to ensure each sample/batch gets evaluated exactly once. This is helpful to make sure benchmarking for research papers is done the right way. Otherwise, in a multi-device setting, samples could occur duplicated when DistributedSampleris used, for eg. with strategy="ddp". It replicates some samples on some devices to make sure all devices have same batch size in case of uneven inputs.

Validation Epoch-level Metrics¶

In the case that you need to make use of all the outputs from each validation_step(), override the on_validation_epoch_end() method. Note that this method is called before on_train_epoch_end().

class LightningTransformer(L.LightningModule): def init(self, vocab_size): super().init() self.model = Transformer(vocab_size=vocab_size) self.validation_step_outputs = []

def validation_step(self, batch, batch_idx):
    x, y = batch
    inputs, target = batch
    output = self.model(inputs, target)
    loss = torch.nn.functional.nll_loss(output, target.view(-1))
    pred = ...
    self.validation_step_outputs.append(pred)
    return pred

def on_validation_epoch_end(self):
    all_preds = torch.stack(self.validation_step_outputs)
    # do something with all preds
    ...
    self.validation_step_outputs.clear()  # free memory

Testing¶

Test Loop¶

The process for enabling a test loop is the same as the process for enabling a validation loop. Please refer to the section above for details. For this you need to override the test_step() method.

The only difference is that the test loop is only called when test() is used.

model = LightningTransformer(vocab_size=dataset.vocab_size) dataloader = DataLoader(dataset) trainer = L.Trainer() trainer.fit(model=model, train_dataloaders=dataloader)

automatically loads the best weights for you

trainer.test(model)

There are two ways to call test():

call after training

trainer = L.Trainer() trainer.fit(model=model, train_dataloaders=dataloader)

automatically auto-loads the best weights from the previous run

trainer.test(dataloaders=test_dataloaders)

or call with pretrained model

model = LightningTransformer.load_from_checkpoint(PATH) dataset = WikiText2() test_dataloader = DataLoader(dataset) trainer = L.Trainer() trainer.test(model, dataloaders=test_dataloader)

Note

WikiText2 is used in a manner that does not create a train, test, val split. This is done for illustrative purposes only. A proper split can be created in lightning.pytorch.core.LightningModule.setup() or lightning.pytorch.core.LightningDataModule.setup().

Note

Inference¶

Prediction Loop¶

By default, the predict_step() method runs theforward() method. In order to customize this behaviour, simply override the predict_step() method.

For the example let’s override predict_step:

class LightningTransformer(L.LightningModule): def init(self, vocab_size): super().init() self.model = Transformer(vocab_size=vocab_size)

def predict_step(self, batch):
    inputs, target = batch
    return self.model(inputs, target)

Under the hood, Lightning does the following (pseudocode):

disable grads + batchnorm + dropout

torch.set_grad_enabled(False) model.eval() all_preds = []

for batch_idx, batch in enumerate(predict_dataloader): pred = model.predict_step(batch, batch_idx) all_preds.append(pred)

There are two ways to call predict():

call after training

trainer = L.Trainer() trainer.fit(model=model, train_dataloaders=dataloader)

automatically auto-loads the best weights from the previous run

predictions = trainer.predict(dataloaders=predict_dataloader)

or call with pretrained model

model = LightningTransformer.load_from_checkpoint(PATH) dataset = WikiText2() test_dataloader = DataLoader(dataset) trainer = L.Trainer() predictions = trainer.predict(model, dataloaders=test_dataloader)

Inference in Research¶

If you want to perform inference with the system, you can add a forward method to the LightningModule.

Note

When using forward, you are responsible to call eval() and use the no_grad() context manager.

class LightningTransformer(L.LightningModule): def init(self, vocab_size): super().init() self.model = Transformer(vocab_size=vocab_size)

def forward(self, batch):
    inputs, target = batch
    return self.model(inputs, target)

def training_step(self, batch, batch_idx):
    inputs, target = batch
    output = self.model(inputs, target)
    loss = torch.nn.functional.nll_loss(output, target.view(-1))
    return loss

def configure_optimizers(self):
    return torch.optim.SGD(self.model.parameters(), lr=0.1)

model = LightningTransformer(vocab_size=dataset.vocab_size)

model.eval() with torch.no_grad(): batch = dataloader.dataset[0] pred = model(batch)

The advantage of adding a forward is that in complex systems, you can do a much more involved inference procedure, such as text generation:

class Seq2Seq(L.LightningModule): def forward(self, x): embeddings = self(x) hidden_states = self.encoder(embeddings) for h in hidden_states: # decode ... return decoded

In the case where you want to scale your inference, you should be usingpredict_step().

class Autoencoder(L.LightningModule): def forward(self, x): return self.decoder(x)

def predict_step(self, batch, batch_idx, dataloader_idx=0):
    # this calls forward
    return self(batch)

data_module = ... model = Autoencoder() trainer = Trainer(accelerator="gpu", devices=2) trainer.predict(model, data_module)

Inference in Production¶

For cases like production, you might want to iterate different models inside a LightningModule.

from torchmetrics.functional import accuracy

class ClassificationTask(L.LightningModule): def init(self, model): super().init() self.model = model

def training_step(self, batch, batch_idx):
    x, y = batch
    y_hat = self.model(x)
    loss = F.cross_entropy(y_hat, y)
    return loss

def validation_step(self, batch, batch_idx):
    loss, acc = self._shared_eval_step(batch, batch_idx)
    metrics = {"val_acc": acc, "val_loss": loss}
    self.log_dict(metrics)
    return metrics

def test_step(self, batch, batch_idx):
    loss, acc = self._shared_eval_step(batch, batch_idx)
    metrics = {"test_acc": acc, "test_loss": loss}
    self.log_dict(metrics)
    return metrics

def _shared_eval_step(self, batch, batch_idx):
    x, y = batch
    y_hat = self.model(x)
    loss = F.cross_entropy(y_hat, y)
    acc = accuracy(y_hat, y)
    return loss, acc

def predict_step(self, batch, batch_idx, dataloader_idx=0):
    x, y = batch
    y_hat = self.model(x)
    return y_hat

def configure_optimizers(self):
    return torch.optim.Adam(self.model.parameters(), lr=0.02)

Then pass in any arbitrary model to be fit with this task

for model in [resnet50(), vgg16(), BidirectionalRNN()]: task = ClassificationTask(model)

trainer = Trainer(accelerator="gpu", devices=2)
trainer.fit(task, train_dataloaders=train_dataloader, val_dataloaders=val_dataloader)

Tasks can be arbitrarily complex such as implementing GAN training, self-supervised or even RL.

class GANTask(L.LightningModule): def init(self, generator, discriminator): super().init() self.generator = generator self.discriminator = discriminator

...

When used like this, the model can be separated from the Task and thus used in production without needing to keep it in a LightningModule.

The following example shows how you can run inference in the Python runtime:

task = ClassificationTask(model) trainer = Trainer(accelerator="gpu", devices=2) trainer.fit(task, train_dataloader, val_dataloader) trainer.save_checkpoint("best_model.ckpt")

use model after training or load weights and drop into the production system

model = ClassificationTask.load_from_checkpoint("best_model.ckpt") x = ... model.eval() with torch.no_grad(): y_hat = model(x)

Check out Inference in Production guide to learn about the possible ways to perform inference in production.

Save Hyperparameters¶

Often times we train many versions of a model. You might share that model or come back to it a few months later at which point it is very useful to know how that model was trained (i.e.: what learning rate, neural network, etc…).

Lightning has a standardized way of saving the information for you in checkpoints and YAML files. The goal here is to improve readability and reproducibility.

save_hyperparameters¶

Use save_hyperparameters() within yourLightningModule’s __init__ method. It will enable Lightning to store all the provided arguments under the self.hparams attribute. These hyperparameters will also be stored within the model checkpoint, which simplifies model re-instantiation after training.

class LitMNIST(L.LightningModule): def init(self, layer_1_dim=128, learning_rate=1e-2): super().init() # call this to save (layer_1_dim=128, learning_rate=1e-4) to the checkpoint self.save_hyperparameters()

    # equivalent
    self.save_hyperparameters("layer_1_dim", "learning_rate")

    # Now possible to access layer_1_dim from hparams
    self.hparams.layer_1_dim

In addition, loggers that support it will automatically log the contents of self.hparams.

Excluding hyperparameters¶

By default, every parameter of the __init__ method will be considered a hyperparameter to the LightningModule. However, sometimes some parameters need to be excluded from saving, for example when they are not serializable. Those parameters should be provided back when reloading the LightningModule. In this case, exclude them explicitly:

class LitMNIST(L.LightningModule): def init(self, loss_fx, generator_network, layer_1_dim=128): super().init() self.layer_1_dim = layer_1_dim self.loss_fx = loss_fx

    # call this to save only (layer_1_dim=128) to the checkpoint
    self.save_hyperparameters("layer_1_dim")

    # equivalent
    self.save_hyperparameters(ignore=["loss_fx", "generator_network"])

load_from_checkpoint¶

LightningModules that have hyperparameters automatically saved withsave_hyperparameters() can conveniently be loaded and instantiated directly from a checkpoint with load_from_checkpoint():

to load specify the other args

model = LitMNIST.load_from_checkpoint(PATH, loss_fx=torch.nn.SomeOtherLoss, generator_network=MyGenerator())

If parameters were excluded, they need to be provided at the time of loading:

the excluded parameters were `loss_fx` and `generator_network`

model = LitMNIST.load_from_checkpoint(PATH, loss_fx=torch.nn.SomeOtherLoss, generator_network=MyGenerator())

Child Modules¶

Research projects tend to test different approaches to the same dataset. This is very easy to do in Lightning with inheritance.

For example, imagine we now want to train an AutoEncoder to use as a feature extractor for images. The only things that change in the LitAutoEncoder model are the init, forward, training, validation and test step.

class Encoder(torch.nn.Module): ...

class Decoder(torch.nn.Module): ...

class AutoEncoder(torch.nn.Module): def init(self): super().init() self.encoder = Encoder() self.decoder = Decoder()

def forward(self, x):
    return self.decoder(self.encoder(x))

class LitAutoEncoder(LightningModule): def init(self, auto_encoder): super().init() self.auto_encoder = auto_encoder self.metric = torch.nn.MSELoss()

def forward(self, x):
    return self.auto_encoder.encoder(x)

def training_step(self, batch, batch_idx):
    x, _ = batch
    x_hat = self.auto_encoder(x)
    loss = self.metric(x, x_hat)
    return loss

def validation_step(self, batch, batch_idx):
    self._shared_eval(batch, batch_idx, "val")

def test_step(self, batch, batch_idx):
    self._shared_eval(batch, batch_idx, "test")

def _shared_eval(self, batch, batch_idx, prefix):
    x, _ = batch
    x_hat = self.auto_encoder(x)
    loss = self.metric(x, x_hat)
    self.log(f"{prefix}_loss", loss)

and we can train this using the Trainer:

auto_encoder = AutoEncoder() lightning_module = LitAutoEncoder(auto_encoder) trainer = Trainer() trainer.fit(lightning_module, train_dataloader, val_dataloader)

And remember that the forward method should define the practical use of a LightningModule. In this case, we want to use the LitAutoEncoder to extract image representations:

some_images = torch.Tensor(32, 1, 28, 28) representations = lightning_module(some_images)

LightningModule API¶

Methods¶

all_gather¶

LightningModule.all_gather(data, group=None, sync_grads=False)[source]

Gather tensors or collections of tensors from multiple processes.

This method needs to be called on all processes and the tensors need to have the same shape across all processes, otherwise your program will stall forever.

Parameters:

data¶ (Union[Tensor, dict, list, tuple]) – int, float, tensor of shape (batch, …), or a (possibly nested) collection thereof.
group¶ (Optional[Any]) – the process group to gather results from. Defaults to all processes (world)
sync_grads¶ (bool) – flag that allows users to synchronize gradients for the all_gather operation

Return type:

Union[Tensor, dict, list, tuple]

Returns:

A tensor of shape (world_size, batch, …), or if the input was a collection the output will also be a collection with tensors of this shape. For the special case where world_size is 1, no additional dimension is added to the tensor(s).

configure_callbacks¶

LightningModule.configure_callbacks()[source]

Configure model-specific callbacks. When the model gets attached, e.g., when .fit() or .test() gets called, the list or a callback returned here will be merged with the list of callbacks passed to the Trainer’scallbacks argument. If a callback returned here has the same type as one or several callbacks already present in the Trainer’s callbacks list, it will take priority and replace them. In addition, Lightning will make sure ModelCheckpoint callbacks run last.

Return type:

Union[Sequence[Callback], Callback]

Returns:

A callback or a list of callbacks which will extend the list of callbacks in the Trainer.

Example:

def configure_callbacks(self): early_stop = EarlyStopping(monitor="val_acc", mode="max") checkpoint = ModelCheckpoint(monitor="val_loss") return [early_stop, checkpoint]

configure_optimizers¶

LightningModule.configure_optimizers()[source]

Choose what optimizers and learning-rate schedulers to use in your optimization. Normally you’d need one. But in the case of GANs or similar you might have multiple. Optimization with multiple optimizers only works in the manual optimization mode.

Return type:

Union[Optimizer, Sequence[Optimizer], tuple[Sequence[Optimizer], Sequence[Union[LRScheduler, ReduceLROnPlateau, LRSchedulerConfig]]], OptimizerConfig, OptimizerLRSchedulerConfig, Sequence[OptimizerConfig], Sequence[OptimizerLRSchedulerConfig], None]

Returns:

Any of these 6 options.

Single optimizer.
List or Tuple of optimizers.
Two lists - The first list has multiple optimizers, and the second has multiple LR schedulers (or multiple lr_scheduler_config).
Dictionary, with an "optimizer" key, and (optionally) a "lr_scheduler"key whose value is a single LR scheduler or lr_scheduler_config.
None - Fit will run without any optimizer.

The lr_scheduler_config is a dictionary which contains the scheduler and its associated configuration. The default configuration is shown below.

lr_scheduler_config = { # REQUIRED: The scheduler instance "scheduler": lr_scheduler, # The unit of the scheduler's step size, could also be 'step'. # 'epoch' updates the scheduler on epoch end whereas 'step' # updates it after a optimizer update. "interval": "epoch", # How many epochs/steps should pass between calls to # scheduler.step(). 1 corresponds to updating the learning # rate after every epoch/step. "frequency": 1, # Metric to monitor for schedulers like ReduceLROnPlateau "monitor": "val_loss", # If set to True, will enforce that the value specified 'monitor' # is available when the scheduler is updated, thus stopping # training if not found. If set to False, it will only produce a warning "strict": True, # If using the LearningRateMonitor callback to monitor the # learning rate progress, this keyword can be used to specify # a custom logged name "name": None, }

When there are schedulers in which the .step() method is conditioned on a value, such as thetorch.optim.lr_scheduler.ReduceLROnPlateau scheduler, Lightning requires that thelr_scheduler_config contains the keyword "monitor" set to the metric name that the scheduler should be conditioned on.

The ReduceLROnPlateau scheduler requires a monitor

def configure_optimizers(self): optimizer = Adam(...) return { "optimizer": optimizer, "lr_scheduler": { "scheduler": ReduceLROnPlateau(optimizer, ...), "monitor": "metric_to_track", "frequency": "indicates how often the metric is updated", # If "monitor" references validation metrics, then "frequency" should be set to a # multiple of "trainer.check_val_every_n_epoch". }, }

In the case of two optimizers, only one using the ReduceLROnPlateau scheduler

def configure_optimizers(self): optimizer1 = Adam(...) optimizer2 = SGD(...) scheduler1 = ReduceLROnPlateau(optimizer1, ...) scheduler2 = LambdaLR(optimizer2, ...) return ( { "optimizer": optimizer1, "lr_scheduler": { "scheduler": scheduler1, "monitor": "metric_to_track", }, }, {"optimizer": optimizer2, "lr_scheduler": scheduler2}, )

Metrics can be made available to monitor by simply logging it usingself.log('metric_to_track', metric_val) in your LightningModule.

Note

Some things to know:

Lightning calls .backward() and .step() automatically in case of automatic optimization.
If a learning rate scheduler is specified in configure_optimizers() with key"interval" (default “epoch”) in the scheduler configuration, Lightning will call the scheduler’s .step() method automatically in case of automatic optimization.
If you use 16-bit precision (precision=16), Lightning will automatically handle the optimizer.
If you use torch.optim.LBFGS, Lightning handles the closure function automatically for you.
If you use multiple optimizers, you will have to switch to ‘manual optimization’ mode and step them yourself.
If you need to control how often the optimizer steps, override the optimizer_step() hook.

forward¶

LightningModule.forward(*args, **kwargs)[source]

Same as torch.nn.Module.forward().

Parameters:

*args¶ (Any) – Whatever you decide to pass into the forward method.
**kwargs¶ (Any) – Keyword arguments are also possible.

Return type:

Any

Returns:

Your model’s output

freeze¶

LightningModule.freeze()[source]

Freeze all params for inference.

Example:

model = MyLightningModule(...) model.freeze()

Return type:

LightningModule — PyTorch Lightning 2.5.1.post0 documentation (original) (raw)

don't do in Lightning

do this instead

or to init a new tensor

Don't do in Lightning...

do this instead

Starter Example¶

Training¶

Training Loop¶

enable gradient calculation

Train Epoch-level Metrics¶

note: in reality, we do this incrementally, instead of keeping all outputs in memory

Train Epoch-level Operations¶

Validation¶

Validation Loop¶

...

Validation Epoch-level Metrics¶

Testing¶

Test Loop¶

automatically loads the best weights for you

call after training

automatically auto-loads the best weights from the previous run

or call with pretrained model

Inference¶

Prediction Loop¶

disable grads + batchnorm + dropout

call after training

automatically auto-loads the best weights from the previous run

or call with pretrained model

Inference in Research¶

Inference in Production¶

use model after training or load weights and drop into the production system

Save Hyperparameters¶

save_hyperparameters¶

Excluding hyperparameters¶

load_from_checkpoint¶

to load specify the other args

the excluded parameters were loss_fx and generator_network

Child Modules¶

LightningModule API¶

Methods¶

all_gather¶

configure_callbacks¶

configure_optimizers¶

The ReduceLROnPlateau scheduler requires a monitor

In the case of two optimizers, only one using the ReduceLROnPlateau scheduler

forward¶

freeze¶

log¶

log_dict¶

lr_schedulers¶

manual_backward¶

optimizers¶

print¶

predict_step¶

save_hyperparameters¶

toggle_optimizer¶

test_step¶

if you have one test dataloader:

if you have multiple test dataloaders:

CASE 1: A single test dataset

CASE 2: multiple test dataloaders

to_onnx¶

to_torchscript¶

training_step¶

Multiple optimizers (e.g.: GANs)

unfreeze¶

untoggle_optimizer¶

validation_step¶

if you have one val dataloader:

if you have multiple val dataloaders:

CASE 1: A single validation dataset

CASE 2: multiple validation dataloaders

Properties¶

current_epoch¶

device¶

global_rank¶

global_step¶

hparams¶

logger¶

the excluded parameters were `loss_fx` and `generator_network`

Set gradients to `None` instead of zero to improve performance (not required on `torch>=2.0.0`).