Manual Optimization Mixin#

Module for the manual optimization mixin class.

class ManualOptimizationMixin[source]

Bases: object

Mixin that handles Lightning manual optimization loops, useful for solvers that require explicit control over optimization steps, such as those with multiple optimizers or custom training loops.

Designed to be used in combination with any solver inheriting from BaseSolver.

training_step(batch, batch_idx)[source]

Solver training step.

Parameters:

  • batch (list[tuple[str, dict]]) – A batch of data. Each element is a tuple containing a condition name and a dictionary of points.

  • batch_idx (int) – The index of the current batch.

Returns:

The loss of the training step.

Return type:

torch.Tensor

on_train_batch_end(outputs, batch, batch_idx)[source]

Keep Lightning’s manual optimization progress counters in sync.

This hook increments the completed optimization-step counter used by Lightning’s manual optimization loop, then delegates to the parent implementation.

Parameters:

  • outputs (torch.Tensor) – The loss of the training step.

  • batch (list[tuple[str, dict]]) – A batch of data. Each element is a tuple containing a condition name and a dictionary of points.

  • batch_idx (int) – The index of the current batch.

Returns:

The result returned by the parent class implementation.

Return type:

Any