def training_step(self, batch, batch_idx):
    utils.update_momentum(self.backbone, self.teacher_backbone, m=0.99)
    utils.update_momentum(self.head, self.teacher_head, m=0.99)
    (views, _, _) = batch
    views = [view.to(self.device) for view in views]
    global_views = views[:2]
    teacher_out = [self.forward_teacher(view) for view in global_views]
    student_out = [self.forward(view) for view in views]
    loss = self.criterion(teacher_out, student_out, epoch=self.current_epoch)
    self.log('train_loss_ssl', loss)
    return loss
def training_step(self, batch, batch_idx):
    utils.update_momentum(self.backbone, self.teacher_backbone, m=0.99)
    utils.update_momentum(self.head, self.teacher_head, m=0.99)
    (views, _, _) = batch
    views = [view.to(self.device) for view in views]
    global_views = views[:2]
    teacher_out = [self.forward_teacher(view) for view in global_views]
    student_out = [self.forward(view) for view in views]
    loss = self.criterion(teacher_out, student_out, epoch=self.current_epoch)
    self.log('train_loss_ssl', loss)
    return loss