def training_step(self, batch, batch_idx):
    ((x0, x1), _, _) = batch
    utils.update_momentum(self.backbone, self.backbone_momentum, 0.99)
    utils.update_momentum(self.projection_head, self.projection_head_momentum, 0.99)

    def step(x0_, x1_):
        (x1_, shuffle) = utils.batch_shuffle(x1_, distributed=distributed)
        x0_ = self.backbone(x0_).flatten(start_dim=1)
        x0_ = self.projection_head(x0_)
        x1_ = self.backbone_momentum(x1_).flatten(start_dim=1)
        x1_ = self.projection_head_momentum(x1_)
        x1_ = utils.batch_unshuffle(x1_, shuffle, distributed=distributed)
        return (x0_, x1_)
    loss_1 = self.criterion(*step(x0, x1))
    loss_2 = self.criterion(*step(x1, x0))
    loss = (0.5 * (loss_1 + loss_2))
    self.log('train_loss_ssl', loss)
    return loss
def training_step(self, batch, batch_idx):
    ((x0, x1), _, _) = batch
    utils.update_momentum(self.backbone, self.backbone_momentum, 0.99)
    utils.update_momentum(self.projection_head, self.projection_head_momentum, 0.99)

    def step(x0_, x1_):
        (x1_, shuffle) = utils.batch_shuffle(x1_, distributed=distributed)
        x0_ = self.backbone(x0_).flatten(start_dim=1)
        x0_ = self.projection_head(x0_)
        x1_ = self.backbone_momentum(x1_).flatten(start_dim=1)
        x1_ = self.projection_head_momentum(x1_)
        x1_ = utils.batch_unshuffle(x1_, shuffle, distributed=distributed)
        return (x0_, x1_)
    loss_1 = self.criterion(*step(x0, x1))
    loss_2 = self.criterion(*step(x1, x0))
    loss = (0.5 * (loss_1 + loss_2))
    self.log('train_loss_ssl', loss)
    return loss