def forward_backward(self, features, targets, optimizer):
        """
        Partial FC forward, which will sample positive weights and part of negative weights,
        then compute logits and get the grad of features.
        """
        total_targets = self.prepare(targets, optimizer)

        if self.world_size > 1:
            total_features = concat_all_gather(features)
        else:
            total_features = features.detach()

        total_features.requires_grad_(True)

        logits = self.forward(total_features)
        logits = self.cls_layer(logits, total_targets)

        # from ipdb import set_trace; set_trace()
        with torch.no_grad():
            max_fc = torch.max(logits, dim=1, keepdim=True)[0]
            if self.world_size > 1:
                dist.all_reduce(max_fc, dist.ReduceOp.MAX)

            # calculate exp(logits) and all-reduce
            logits_exp = torch.exp(logits - max_fc)
            logits_sum_exp = logits_exp.sum(dim=1, keepdim=True)

            if self.world_size > 1:
                dist.all_reduce(logits_sum_exp, dist.ReduceOp.SUM)

            # calculate prob
            logits_exp.div_(logits_sum_exp)

            # get one-hot
            grad = logits_exp
            index = torch.where(total_targets != -1)[0]
            one_hot = torch.zeros(size=[index.size()[0], grad.size()[1]], device=grad.device)
            one_hot.scatter_(1, total_targets[index, None], 1)

            # calculate loss
            loss = torch.zeros(grad.size()[0], 1, device=grad.device)
            loss[index] = grad[index].gather(1, total_targets[index, None])
            if self.world_size > 1:
                dist.all_reduce(loss, dist.ReduceOp.SUM)
            loss_v = loss.clamp_min_(1e-30).log_().mean() * (-1)

            # calculate grad
            grad[index] -= one_hot
            grad.div_(logits.size(0))

        logits.backward(grad)
        if total_features.grad is not None:
            total_features.grad.detach_()
        x_grad: torch.Tensor = torch.zeros_like(features)
        # feature gradient all-reduce
        if self.world_size > 1:
            dist.reduce_scatter(x_grad, list(total_features.grad.chunk(self.world_size, dim=0)))
        else:
            x_grad = total_features.grad
        x_grad = x_grad * self.world_size
        # backward backbone
        return x_grad, loss_v

        Partial FC forward, which will sample positive weights and part of negative weights,
        then compute logits and get the grad of features.
        
projects/FastFace/fastface/modeling/partial_fc.py