def forward(self, x):
    f = self.backbone(x).flatten(start_dim=1)
    z = self.projection_head(f)
    p = self.prediction_head(z)
    z = z.detach()
    return (z, p)
def forward(self, x):
    f = self.backbone(x).flatten(start_dim=1)
    z = self.projection_head(f)
    p = self.prediction_head(z)
    z = z.detach()
    return (z, p)