def __init__(
        self,
        *,
        model: Module,
        optimizer: Optimizer,
        criterion: CriterionType = CrossEntropyLoss(),
        train_mb_size: int = 1,
        train_epochs: int = 1,
        eval_mb_size: Optional[int] = None,
        device=None,
        plugins: Optional[List[SupervisedPlugin]] = None,
        evaluator: EvaluationPlugin = default_evaluator,
        eval_every=-1,
        ae_train_mb_size=1,
        ae_train_epochs=2,
        ae_latent_dim=100,
        ae_lr=1e-3,
        temp=2,
        rel_thresh=0.85,
        **base_kwargs
    ):
        """Init.

        :param model: PyTorch model.
        :param optimizer: PyTorch optimizer.
        :param criterion: loss function.
        :param train_mb_size: mini-batch size for training.
        :param train_epochs: number of training epochs.
        :param eval_mb_size: mini-batch size for eval.
        :param device: PyTorch device where the model will be allocated.
        :param plugins: (optional) list of StrategyPlugins.
        :param evaluator: (optional) instance of EvaluationPlugin for logging
            and metric computations. None to remove logging.
        :param eval_every: the frequency of the calls to `eval` inside the
            training loop. -1 disables the evaluation. 0 means `eval` is called
            only at the end of the learning experience. Values >0 mean that
            `eval` is called every `eval_every` epochs and at the end of the
            learning experience.
        :param base_kwargs: any additional
        :param ae_train_mb_size: mini-batch size for training of the autoencoder
        :param ae_train_epochs: number of training epochs for the autoencoder
        :param ae_lr: the learning rate for the autoencoder training
        using vanilla SGD
        :param temp: the temperature hyperparameter when selecting the
        expert during the forward method
            :class:`~avalanche.training.BaseTemplate` constructor arguments.
        """

        expertgate = _ExpertGatePlugin()

        if plugins is None:
            plugins = [expertgate]
        else:
            plugins += [expertgate]

        self.ae_train_mb_size = ae_train_mb_size
        self.ae_train_epochs = ae_train_epochs
        self.ae_lr = ae_lr
        self.ae_latent_dim = ae_latent_dim
        self.rel_thresh = rel_thresh

        warnings.warn(
            "This strategy is currently in the alpha stage and we are still "
            "working to reproduce the original paper's results. You can find "
            "the code to reproduce the experiments at "
            "github.com/continualAI/continual-learning-baselines"
        )

        super().__init__(
            model=model,
            optimizer=optimizer,
            criterion=criterion,
            train_mb_size=train_mb_size,
            train_epochs=train_epochs,
            eval_mb_size=eval_mb_size,
            device=device,
            plugins=plugins,
            evaluator=evaluator,
            eval_every=eval_every,
            **base_kwargs
        )

        # Check that the model has the correct architecture.
        assert isinstance(
            self.model, ExpertGate
        ), "ExpertGateStrategy requires an ExpertGate model."
        self.model.temp = temp
Init.

        :param model: PyTorch model.
        :param optimizer: PyTorch optimizer.
        :param criterion: loss function.
        :param train_mb_size: mini-batch size for training.
        :param train_epochs: number of training epochs.
        :param eval_mb_size: mini-batch size for eval.
        :param device: PyTorch device where the model will be allocated.
        :param plugins: (optional) list of StrategyPlugins.
        :param evaluator: (optional) instance of EvaluationPlugin for logging
            and metric computations. None to remove logging.
        :param eval_every: the frequency of the calls to `eval` inside the
            training loop. -1 disables the evaluation. 0 means `eval` is called
            only at the end of the learning experience. Values >0 mean that
            `eval` is called every `eval_every` epochs and at the end of the
            learning experience.
        :param base_kwargs: any additional
        :param ae_train_mb_size: mini-batch size for training of the autoencoder
        :param ae_train_epochs: number of training epochs for the autoencoder
        :param ae_lr: the learning rate for the autoencoder training
        using vanilla SGD
        :param temp: the temperature hyperparameter when selecting the
        expert during the forward method
            :class:`~avalanche.training.BaseTemplate` constructor arguments.
        
avalanche/training/supervised/expert_gate.py