class_docstr: EfficientAd model.
  
  Args:
      teacher_out_channels (int): number of convolution output channels of the pre-trained teacher model
      pretrained_models_dir (str): path to the pretrained model weights
      input_size (tuple): size of input images
      model_size (str): size of student and teacher model
      padding (bool): use padding in convoluional layers
      pad_maps (bool): relevant if padding is set to False. In this case, pad_maps = True pads the
          output anomaly maps so that their size matches the size in the padding = True case.
      device (str): which device the model should be loaded on
docstr: Prediction by EfficientAd models.
  
  Args:
      normalize: Normalize anomaly maps or not.
      batch (Tensor): Input images.
  
  Returns:
      Tensor: Predictions
code: def forward(
      self,
      batch: Tensor,
      batch_imagenet: Tensor = None,
      normalize: bool = True,
  ) -> Tensor | dict:
      """Prediction by EfficientAd models.
  
      Args:
          normalize: Normalize anomaly maps or not.
          batch (Tensor): Input images.
  
      Returns:
          Tensor: Predictions
      """
      with torch.no_grad():
          teacher_output = self.teacher(batch)
          if self.is_set(self.mean_std):
              teacher_output = (
                  teacher_output - self.mean_std["mean"]
              ) / self.mean_std["std"]
  
      student_output = self.student(batch)
      distance_st = torch.pow(
          teacher_output
          - student_output[:, : self.teacher_out_channels, :, :],
          2,
      )
  
      if self.training:
          # Student loss
          distance_st = reduce_tensor_elems(distance_st)
          d_hard = torch.quantile(distance_st, 0.999)
          loss_hard = torch.mean(distance_st[distance_st >= d_hard])
          student_output_penalty = self.student(batch_imagenet)[
              :, : self.teacher_out_channels, :, :
          ]
          loss_penalty = torch.mean(student_output_penalty**2)
          loss_st = loss_hard + loss_penalty
  
          # Autoencoder and Student AE Loss
          aug_img = self.choose_random_aug_image(batch)
          ae_output_aug = self.ae(aug_img)
  
          with torch.no_grad():
              teacher_output_aug = self.teacher(aug_img)
              if self.is_set(self.mean_std):
                  teacher_output_aug = (
                      teacher_output_aug - self.mean_std["mean"]
                  ) / self.mean_std["std"]
  
          student_output_ae_aug = self.student(aug_img)[
              :, self.teacher_out_channels :, :, :
          ]
  
          distance_ae = torch.pow(teacher_output_aug - ae_output_aug, 2)
          distance_stae = torch.pow(
              ae_output_aug - student_output_ae_aug, 2
          )
  
          loss_ae = torch.mean(distance_ae)
          loss_stae = torch.mean(distance_stae)
          return (loss_st, loss_ae, loss_stae)
  
      else:
          with torch.no_grad():
              ae_output = self.ae(batch)
  
          map_st = torch.mean(distance_st, dim=1, keepdim=True)
          map_stae = torch.mean(
              (
                  ae_output
                  - student_output[:, self.teacher_out_channels :]
              )
              ** 2,
              dim=1,
              keepdim=True,
          )
  
          if self.pad_maps:
              map_st = F.pad(map_st, (4, 4, 4, 4))
              map_stae = F.pad(map_stae, (4, 4, 4, 4))
          map_st = F.interpolate(
              map_st,
              size=(self.input_size[0], self.input_size[1]),
              mode="bilinear",
          )
          map_stae = F.interpolate(
              map_stae,
              size=(self.input_size[0], self.input_size[1]),
              mode="bilinear",
          )
  
          if self.is_set(self.quantiles) and normalize:
              map_st = (
                  0.1
                  * (map_st - self.quantiles["qa_st"])
                  / (self.quantiles["qb_st"] - self.quantiles["qa_st"])
              )
              map_stae = (
                  0.1
                  * (map_stae - self.quantiles["qa_ae"])
                  / (self.quantiles["qb_ae"] - self.quantiles["qa_ae"])
              )
  
          map_combined = 0.5 * map_st + 0.5 * map_stae
          return {
              "anomaly_map": map_combined,
              "map_st": map_st,
              "map_ae": map_stae,
          }
  
imports: from __future__ import annotations
  import logging
  import random
  from enum import Enum
  import torch
  import torch.nn.functional
  from torch import Tensor, nn
  from torchvision import transforms
label_desc: Use image node_id 4 for calling the EfficientAdModel obj's underlying member method code with example usage: obj.forward(batch, batch_imagenet, normalize) and returns: loss_st,  loss_ae,  loss_stae, dict