class_docstr: Fully Convolutional Cross-Scale-Flows for Image-based Defect Detection.
  
  Args:
      input_size (tuple[int, int]): Size of the model input.
      n_coupling_blocks (int): Number of coupling blocks in the model.
      cross_conv_hidden_channels (int): Number of hidden channels in the cross convolution.
      clamp (int): Clamp value for glow layer.
      num_channels (int): Number of channels in the model.
class_docstr_tok: ["Fully","Convolutional","Cross-Scale-Flows","for","Image-based","Defect","Detection",".","Args",":","input_size","(","tuple","[","int",",","int","]",")",":","Size","of","the","model","input",".","n_coupling_blocks","(","int",")",":","Number","of","coupling","blocks","in","the","model",".","cross_conv_hidden_channels","(","int",")",":","Number","of","hidden","channels","in","the","cross","convolution",".","clamp","(","int",")",":","Clamp","value","for","glow","layer",".","num_channels","(","int",")",":","Number","of","channels","in","the","model","."]
docstr: Training Step of CS-Flow.
  
  Args:
      batch (dict[str, str | Tensor]): Input batch
      _: Index of the batch.
  
  Returns:
      Loss value
docstr_tok: ["Training","Step","of","CS-Flow",".","Args",":","batch","(","dict","[","str",",","str","|","Tensor","]",")",":","Input","batch","_",":","Index","of","the","batch",".","Returns",":","Loss","value"]
code: def training_step(
      self, batch: dict[str, str | Tensor], *args, **kwargs
  ) -> STEP_OUTPUT:
      """Training Step of CS-Flow.
  
      Args:
          batch (dict[str, str | Tensor]): Input batch
          _: Index of the batch.
  
      Returns:
          Loss value
      """
      del args, kwargs  # These variables are not used.
  
      self.model.feature_extractor.eval()
      z_dist, jacobians = self.model(batch["image"])
      loss = self.loss(z_dist, jacobians)
      self.log(
          "train_loss",
          loss.item(),
          on_epoch=True,
          prog_bar=True,
          logger=True,
      )
      return {"loss": loss}
  
imports: from __future__ import annotations
  import logging
  import torch
  from omegaconf import DictConfig, ListConfig
  from pytorch_lightning.callbacks import Callback, EarlyStopping
  from pytorch_lightning.utilities.types import STEP_OUTPUT
  from torch import Tensor
  from anomalib.models.components import AnomalyModule
  from .loss import CsFlowLoss
  from .torch_model import CsFlowModel
label_desc: Use image node_id 2 for calling the Csflow obj's underlying member method code with example usage: obj.training_step(batch) and returns: dict