def __init__(
        self,
        shape,
        device,
        arch="alexnet",
        pretrained_flag=True,
        output_layer_name="features",
    ):
        """
        :param shape: shape of the input layer
        :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__()

        # Store variables
        self.shape = shape
        self.arch = arch
        self.pretrained_flag = pretrained_flag
        self.device = device

        # Dictionary for autoencoders
        # {task, autoencoder}
        self.autoencoder_dict = nn.ModuleDict()

        # Dictionary for experts
        # {task, expert}
        self.expert_dict = nn.ModuleDict()

        # Initialize an expert with pretrained AlexNet
        self.expert = (
            models.__dict__[arch](
                weights=(
                    "AlexNet_Weights.IMAGENET1K_V1"
                    if pretrained_flag
                    else "AlexNet_Weights.NONE"
                )
            )
            .to(device)
            .eval()
        )

        :param shape: shape of the input layer
        :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