CondAEModel.training_step

CondAEModel.training_step(batch: Tuple[torch.Tensor, torch.Tensor], batch_idx: int) → float[source]

Perform a single training step.

Parameters:

  • batch (Tuple[torch.Tensor, torch.Tensor]) – A tuple containing the data for a single batch.

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

Returns:

float – The training loss.