Как оптимизировать создание матрицы смежности с весами расстояний между пикселями граничных блоков в факеле Тензор?

Question

У меня есть тензор граничных блоков над изображениями с формой (количество блоков, 4) для каждого изображения. Коробки представлены в виде [x1, y1, x2, y2] или двух противоположных углов. Во время предварительной обработки я хотел бы создать полностью связанный график, представляющий расстояния в пикселях между всеми полями. У меня есть следующий код, но он очень медленный для целей глубокого обучения:

import torch 
import numpy as np



def calc_distance(box1, box2):
    mid_point_1 = torch.Tensor([ box1[2] - box1[0] , box1[3]- box1[1]])
    mid_point_2 = torch.Tensor([ box2[2] - box2[0] , box2[3] - box2[1]])
    return  np.linalg.norm (mid_point_1 - mid_point_2)


def get_adj_mat(image_boxes:  torch.FloatTensor):
    n_nodes = image_boxes.shape[0]
    A = torch.eye(n_nodes)

    for idx1, box1 in enumerate(image_boxes):
        for idx2 in range(idx1,image_boxes.shape[0]):
            if idx1==idx2:
                continue
            box2 = image_boxes[idx2]
            dist = calc_distance(box1,box2)
            weight= np.exp(-dist/100)
            A[idx1][idx2] = weight
            A[idx2][idx1]= weight
    return torch.FloatTensor(A)

Итак, вопрос в том, есть ли здесь возможность использовать параллельную обработку? Или использовать графический процессор?

Спасибо за помощь.