def get_contrastive_model(
    in_size: int,
    in_channels: int,
    feature_dim: int,
    encoder_dim: int = 512,
    hidden_dim: int = 512,
) -> ContrastiveModel:
    """Init contrastive model based on parsed parametrs.

    Args:
        in_size: size of an image (in_size x in_size)
        in_channels: number of channels in an image
        feature_dim: dimensinality of contrative projection
        encoder_dim: dimensinality of encoder output
        hidden_dim: dimensinality of encoder-contrative projection

    Returns:
        ContrstiveModel instance
    """
    try:
        encoder = resnet9(
            in_size=in_size, in_channels=in_channels, out_features=encoder_dim
        )
    except:
        encoder = resnet_mnist(
            in_size=in_size, in_channels=in_channels, out_features=encoder_dim
        )
    projection_head = nn.Sequential(
        nn.Linear(encoder_dim, hidden_dim, bias=False),
        nn.ReLU(inplace=True),
        nn.Linear(hidden_dim, feature_dim, bias=True),
    )
    model = ContrastiveModel(projection_head, encoder)
    return model
Init contrastive model based on parsed parametrs.

    Args:
        in_size: size of an image (in_size x in_size)
        in_channels: number of channels in an image
        feature_dim: dimensinality of contrative projection
        encoder_dim: dimensinality of encoder output
        hidden_dim: dimensinality of encoder-contrative projection

    Returns:
        ContrstiveModel instance
    
examples/self_supervised/src/common.py