class_docstr: PL Lightning Module for the UFLOW algorithm.
class_docstr_tok: ["PL","Lightning","Module","for","the","UFLOW","algorithm","."]
code: def training_step(
      self, batch: dict[str, str | Tensor], *args, **kwargs
  ) -> STEP_OUTPUT:
      z, ljd = self.model(batch["image"])
      loss = self.loss(z, ljd)
      self.log_dict(
          {"loss": loss},
          on_step=True,
          on_epoch=False,
          prog_bar=False,
          logger=True,
      )
      return {"loss": loss}
  
file_path: anomalib/src/anomalib/models/uflow/lightning_model.py
imports: from __future__ import annotations
  import torch
  from omegaconf import DictConfig, ListConfig
  from pytorch_lightning.callbacks import EarlyStopping
  from pytorch_lightning.utilities.types import STEP_OUTPUT
  from torch import Tensor
  from anomalib.models.components import AnomalyModule
  from anomalib.models.uflow.loss import UFlowLoss
  from anomalib.models.uflow.torch_model import UflowModel
label_desc: Use image node_id 2 for calling the Uflow obj's underlying member method code with example usage: obj.training_step(batch) and returns: dict