def training_step(self, batch: Batch, *args, **kwargs) -> STEP_OUTPUT:
        """Perform a training step of CS-Flow model.

        Args:
            batch (Batch): Input batch containing images and targets
            *args: Additional positional arguments (unused)
            **kwargs: Additional keyword arguments (unused)

        Returns:
            STEP_OUTPUT: Dictionary containing the loss value

        Example:
            >>> batch = Batch(image=torch.randn(32, 3, 256, 256))
            >>> model = Csflow()
            >>> output = model.training_step(batch)
            >>> output["loss"]
            tensor(...)
        """
        del args, kwargs  # These variables are not used.

        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}
Perform a training step of CS-Flow model.

        Args:
            batch (Batch): Input batch containing images and targets
            *args: Additional positional arguments (unused)
            **kwargs: Additional keyword arguments (unused)

        Returns:
            STEP_OUTPUT: Dictionary containing the loss value

        Example:
            >>> batch = Batch(image=torch.randn(32, 3, 256, 256))
            >>> model = Csflow()
            >>> output = model.training_step(batch)
            >>> output["loss"]
            tensor(...)
        
src/anomalib/models/image/csflow/lightning_model.py