def forward(self, x):
    x = self.backbone(x).flatten(start_dim=1)
    x = self.projection_head(x)
    x = nn.functional.normalize(x, dim=1, p=2)
    return self.prototypes(x)
def forward(self, x):
    x = self.backbone(x).flatten(start_dim=1)
    x = self.projection_head(x)
    x = nn.functional.normalize(x, dim=1, p=2)
    return self.prototypes(x)