def __init__(
        self,
        shape,
        latent_dim,
        device,
        arch="alexnet",
        pretrained_flag=True,
        output_layer_name="features",
    ):
        """
        :param shape: shape of the input layer
        :param latent_dim: size of the autoencoder's latent dimension
        :param device: gpu or cpu
        :param arch: the architecture to use from torchvision.models,
        defaults to "alexnet"
        :param pretrained_flag: determines if torchvision model is pre-trained,
        defaults to True
        :param output_layer_name: output layer of the feature backbone,
        defaults to "features"
        """

        super().__init__()

        # Select pretrained AlexNet for preprocessing input
        base_template = models.__dict__[arch](
            weights=(
                "AlexNet_Weights.IMAGENET1K_V1"
                if pretrained_flag
                else "AlexNet_Weights.NONE"
            )
        ).to(device)

        self.feature_module = FeatureExtractorBackbone(base_template, output_layer_name)

        self.feature_module.to(device)

        self.shape = shape
        self.device = device

        # Freeze the feature module
        for param in self.feature_module.parameters():
            param.requires_grad = False

        # Flatten input
        # Encoder Linear -> ReLU
        flattened_size = torch.Size(shape).numel()
        self.encoder = nn.Sequential(
            Flatten(), nn.Linear(flattened_size, latent_dim), nn.ReLU()
        ).to(device)

        # Decoder Linear -> Sigmoid
        self.decoder = nn.Sequential(
            nn.Linear(latent_dim, flattened_size), nn.Sigmoid()
        ).to(device)

        :param shape: shape of the input layer
        :param latent_dim: size of the autoencoder's latent dimension
        :param device: gpu or cpu
        :param arch: the architecture to use from torchvision.models,
        defaults to "alexnet"
        :param pretrained_flag: determines if torchvision model is pre-trained,
        defaults to True
        :param output_layer_name: output layer of the feature backbone,
        defaults to "features"
        
avalanche/models/expert_gate.py