def training_step(self, batch: Batch, *args, **kwargs) -> dict[str, torch.Tensor]:
        """Perform training step.

        Computes student, autoencoder and combined losses using both the input
        batch and a batch from ImageNette.

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

        Returns:
            dict[str, torch.Tensor]: Dictionary containing total loss
        """
        del args, kwargs  # These variables are not used.

        try:
            # infinite dataloader; [0] getting the image not the label
            batch_imagenet = next(self.imagenet_iterator)[0].to(self.device)
        except StopIteration:
            self.imagenet_iterator = iter(self.imagenet_loader)
            batch_imagenet = next(self.imagenet_iterator)[0].to(self.device)

        loss_st, loss_ae, loss_stae = self.model(batch=batch.image, batch_imagenet=batch_imagenet)

        loss = loss_st + loss_ae + loss_stae
        self.log("train_st", loss_st.item(), on_epoch=True, prog_bar=True, logger=True)
        self.log("train_ae", loss_ae.item(), on_epoch=True, prog_bar=True, logger=True)
        self.log("train_stae", loss_stae.item(), on_epoch=True, prog_bar=True, logger=True)
        self.log("train_loss", loss.item(), on_epoch=True, prog_bar=True, logger=True)
        return {"loss": loss}
Perform training step.

        Computes student, autoencoder and combined losses using both the input
        batch and a batch from ImageNette.

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

        Returns:
            dict[str, torch.Tensor]: Dictionary containing total loss
        
src/anomalib/models/image/efficient_ad/lightning_model.py