Вычисление потери функции предсказаний с помощью pytorch

У меня есть сверточная нейронная сеть, которая предсказывает 3 величины: Ux, Uy и P. Это скорость x, скорость y и поле давления. Все они представляют собой 2D-массивы размером [100,60], а размер моего пакета - 10.

Я хочу вычислить потери и обновить сеть, вычислив CURL предсказанной скорости с CURL целевой скорости. У меня есть функция, которая делает это: v = curl (Ux_pred, Uy_pred). Учитывая прогнозируемые Ux и Uy, я хочу вычислить потери, сравнив их с имеющимися у меня целевыми точками: true_curl = curl (Ux_true, Uy_true) - я уже вычислил истинное значение curl и добавил это к моим данным Y, как четвертый канал.

Однако я хочу, чтобы моя сеть предсказывала только Ux, Uy и P. Я хочу, чтобы мои параметры NN обновлялись на основе ПОТЕРИ завитков, чтобы повысить точность Ux и Uy. Потеря локона должна выражаться в единицах Ux и Uy. Я пытался сделать это с помощью автограда Pytorch и уже читал много похожих вопросов, но я просто не могу заставить его работать. Это мой код на данный момент:

        print("pred_Curl shape:", np.shape(pred_curl))
        print("pred_Ux shape:", np.shape(pred[:,0,:,:]))
        print("pred_Uy shape:", np.shape(pred[:,1,:,:]))
        true_curl = torch.from_numpy(y[:,3,:,:]) # not sure where to use the true curl?

        pred_curl = Variable(pred_curl, requires_grad=True)
        
        pred_ux = pred[:,0,:,:]
        pred_uy = pred[:,1,:,:]

        pred_ux = Variable(pred_ux, requires_grad=True)
        pred_uy = Variable(pred_uy, requires_grad=True)

        grad_tensor = torch.autograd.grad(outputs=pred_curl, inputs=(pred_ux, pred_uy), 
                       grad_outputs=torch.ones_like(pred_curl), 
                       retain_graph=True,create_graph=True)

        loss = torch.sum(grad_tensor)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

Это дает следующий результат:

pred_Curl shape: torch.Size([10, 100, 60])
pred_Ux shape: torch.Size([10, 100, 60])
pred_Uy shape: torch.Size([10, 100, 60])

RuntimeError: One of the differentiated Tensors appears to not have been used in the graph. 
Set allow_unused=True if this is the desired behavior.

Любая помощь будет оценена по достоинству!

Изменить: вот моя функция завитка:

    def discrete_curl(self,x,y,new_arr):
            for m in range(100):
                for n in range(60):
                    if n <= 58:
                        if m <= 98:
                            if x[m,n] != 0 and y[m,n] != 0:
                                new_arr[m,n] = ((y[m+1,n] - y[m-1,n]) / 2*1) - ((x[m,n+1] - x[m,n-1]) / 2*1)
            return new_arr 

Где x и y - это Ux ad Uy, а new_arr - это результат curl.