Problems with custom loss functions

[cyhltsrjyyc]

I want to use a custom loss function which needs other inputs, the loss function is showed as follow:

class CustomLoss(nn.Module):
def init(self, teacher_preds_2, teacher_preds_3, teacher_preds_4):
super(CustomLoss, self).init()
self.teacher_preds_2 = torch.tensor(teacher_preds_2, requires_grad=True)
self.teacher_preds_3 = torch.tensor(teacher_preds_3, requires_grad=True)
self.teacher_preds_4 = torch.tensor(teacher_preds_4, requires_grad=True)

def forward(self, y_true, y_pred):
    y_true = torch.tensor(y_true, requires_grad=True)
    y_pred = torch.tensor(y_pred, requires_grad=True)

    main_loss = torch.mean((y_true - y_pred) ** 2)

    distill_loss_2 = sup.get_R2(np.array(y_pred).reshape(1, -1), np.array(self.teacher_preds_2).reshape(1, -1))
    distill_loss_3 = sup.get_R2(np.array(y_pred).reshape(1, -1), np.array(self.teacher_preds_3).reshape(1, -1))
    distill_loss_4 = sup.get_R2(np.array(y_pred).reshape(1, -1), np.array(self.teacher_preds_4).reshape(1, -1))

    distill_loss = (distill_loss_2 + distill_loss_3 * 2 + distill_loss_4 * 7) / 10

    total_loss = main_loss + 0.01 * distill_loss
    return total_loss

The "teacher_preds_2, teacher_preds_3, teacher_preds_4" are three fixed tensors with the same size as the training sample. How can I get the batch of the three tensors and calculate the R2 between them and y_pred during fit process?