def __init__(self, dataloader_kNN, num_classes):
    super().__init__(dataloader_kNN, num_classes)
    resnet = lightly.models.ResNetGenerator('resnet-18')
    self.backbone = nn.Sequential(*list(resnet.children())[:(- 1)], nn.AdaptiveAvgPool2d(1))
    self.projection_head = heads.SwaVProjectionHead(512, 512, 128)
    self.prototypes = heads.SwaVPrototypes(128, 512)
    self.criterion = lightly.loss.SwaVLoss(sinkhorn_gather_distributed=gather_distributed)
def __init__(self, dataloader_kNN, num_classes):
    super().__init__(dataloader_kNN, num_classes)
    resnet = lightly.models.ResNetGenerator('resnet-18')
    self.backbone = nn.Sequential(*list(resnet.children())[:(- 1)], nn.AdaptiveAvgPool2d(1))
    self.projection_head = heads.SwaVProjectionHead(512, 512, 128)
    self.prototypes = heads.SwaVPrototypes(128, 512)
    self.criterion = lightly.loss.SwaVLoss(sinkhorn_gather_distributed=gather_distributed)