N = len(instance_predictions)
        K = 14.7 # average size of groups

        instance_predictions_matrix = \
            instance_predictions.reshape(len(instance_predictions), -1).repeat(
                1, len(instance_predictions)).view(len(instance_predictions), len(instance_predictions))
        squared_diff_btw_predictions = (instance_predictions - instance_predictions_matrix)**2

        first_term = torch.mul(similarities_btw_instances, squared_diff_btw_predictions)
        first_term_loss = torch.sum(first_term)  # requires grad = True


        second_term = (predictions - labels)**2

        second_term_loss = torch.sum(second_term)

        first_loss = 1/(N**2) * first_term_loss
        second_loss = self.trade_off/K * second_term_loss

        loss = first_loss.add(second_loss)
        
        return loss


Here’s my updated loss function, do you think this still breaks computation graph? Or how can I check which operations are breaking computation graph?