def forward(z_dist: list[torch.Tensor], jacobians: torch.Tensor) -> torch.Tensor:
        """Compute the CS-Flow loss.

        Args:
            z_dist (list[torch.Tensor]): List of latent space tensors from each
                scale of the normalizing flow. Each tensor has shape
                ``(batch_size, channels, height, width)``.
            jacobians (torch.Tensor): Log-determinant of the Jacobian matrices
                from the normalizing flows. Shape: ``(batch_size,)``.

        Returns:
            torch.Tensor: Scalar loss value averaged over the batch.

        Example:
            >>> z_dist = [torch.randn(2, 64, 32, 32) for _ in range(3)]
            >>> jacobians = torch.randn(2)
            >>> loss = CsFlowLoss.forward(z_dist, jacobians)
        """
        concatenated = torch.cat([z_dist[i].reshape(z_dist[i].shape[0], -1) for i in range(len(z_dist))], dim=1)
        return torch.mean(0.5 * torch.sum(concatenated**2, dim=(1,)) - jacobians) / concatenated.shape[1]
Compute the CS-Flow loss.

        Args:
            z_dist (list[torch.Tensor]): List of latent space tensors from each
                scale of the normalizing flow. Each tensor has shape
                ``(batch_size, channels, height, width)``.
            jacobians (torch.Tensor): Log-determinant of the Jacobian matrices
                from the normalizing flows. Shape: ``(batch_size,)``.

        Returns:
            torch.Tensor: Scalar loss value averaged over the batch.

        Example:
            >>> z_dist = [torch.randn(2, 64, 32, 32) for _ in range(3)]
            >>> jacobians = torch.randn(2)
            >>> loss = CsFlowLoss.forward(z_dist, jacobians)
        
src/anomalib/models/image/csflow/loss.py