def __init__(
        self, num_classes, arch, device, pretrained_flag, provided_template=None
    ):
        """
        :param num_classes: number of classes this expert model will classify
        :param arch: the architecture to use from torchvision.models
        :param device: gpu or cpu
        :param pretrained_flag: determines if torchvision model is pre-trained
        :param provided_template: the expert model to copy the backbone from,
        defaults to None
        """
        super().__init__()

        self.device = device
        self.num_classes = num_classes

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

        # Set the feature module from provided template
        if provided_template is None:
            self.feature_module = deepcopy(base_template._modules["features"])

        # Use base template if nothing provided
        else:
            self.feature_module = deepcopy(provided_template.feature_module)

        # Set avgpool layer
        self.avg_pool = deepcopy(base_template._modules["avgpool"])

        # Flattener
        self.flatten = Flatten()

        # Classifier module
        self.classifier_module = deepcopy(base_template._modules["classifier"])

        # Customize final layer for  the number of classes in the data
        original_classifier_input_dim = self.classifier_module[-1].in_features
        self.classifier_module[-1] = nn.Linear(
            original_classifier_input_dim, self.num_classes
        )

        for param in self.parameters():
            param.requires_grad = True

        :param num_classes: number of classes this expert model will classify
        :param arch: the architecture to use from torchvision.models
        :param device: gpu or cpu
        :param pretrained_flag: determines if torchvision model is pre-trained
        :param provided_template: the expert model to copy the backbone from,
        defaults to None
        
avalanche/models/expert_gate.py