def __init__(
        self,
        arch="resnet18",
        output_layer_name="layer4.1",
        imagenet_pretrained=True,
        device: Union[str, torch.device] = "cpu",
    ):
        """Init.

        :param arch: backbone architecture. Default is resnet-18, but others
            can be used by modifying layer for
            feature extraction in ``self.feature_extraction_wrapper``.
        :param imagenet_pretrained: True if initializing backbone with imagenet
            pre-trained weights else False
        :param output_layer_name: name of the layer from feature extractor
        :param device: cpu, gpu or other device
        """

        super(SLDAResNetModel, self).__init__()

        feat_extractor = (
            models.__dict__[arch](pretrained=imagenet_pretrained).to(device).eval()
        )
        self.feature_extraction_wrapper = FeatureExtractorBackbone(
            feat_extractor, output_layer_name
        ).eval()
Init.

        :param arch: backbone architecture. Default is resnet-18, but others
            can be used by modifying layer for
            feature extraction in ``self.feature_extraction_wrapper``.
        :param imagenet_pretrained: True if initializing backbone with imagenet
            pre-trained weights else False
        :param output_layer_name: name of the layer from feature extractor
        :param device: cpu, gpu or other device
        
avalanche/models/slda_resnet.py