def training_step(self, batch: Batch, *args, **kwargs) -> STEP_OUTPUT:
        """Perform the training step of FRE.

        For each batch, features are extracted from the CNN backbone and
        reconstructed using the tied autoencoder. The loss is computed as the MSE
        between original and reconstructed features.

        Args:
            batch (Batch): Input batch containing images and labels.
            args: Additional arguments (unused).
            kwargs: Additional keyword arguments (unused).

        Returns:
            STEP_OUTPUT: Dictionary containing the loss value.
        """
        del args, kwargs  # These variables are not used.
        features_in, features_out, _ = self.model.get_features(batch.image)
        loss = self.loss_fn(features_in, features_out)
        self.log("train_loss", loss.item(), on_epoch=True, prog_bar=True, logger=True)
        return {"loss": loss}
Perform the training step of FRE.

        For each batch, features are extracted from the CNN backbone and
        reconstructed using the tied autoencoder. The loss is computed as the MSE
        between original and reconstructed features.

        Args:
            batch (Batch): Input batch containing images and labels.
            args: Additional arguments (unused).
            kwargs: Additional keyword arguments (unused).

        Returns:
            STEP_OUTPUT: Dictionary containing the loss value.
        
src/anomalib/models/image/fre/lightning_model.py