Ошибка вставки тензоров факела в heapq с использованием дублирующих приоритетов

Question

Как избежать RuntimeError: bool value of Tensor with more than one value is ambiguous в этом коде?

import torch
import heapq

h = []
heapq.heappush(h, (1, torch.Tensor([[1,2]])))
heapq.heappush(h, (1, torch.Tensor([[3,4]])))

Это происходит потому, что сравнение кортежей сравнивает вторые элементы, когда первые равны

Answer 1

Необходимо предотвратить попытки heapq сравнивать вторые элементы кортежа, когда он находит дублирующиеся приоритеты и ему нужно только переопределить оператор < для моих элементов.

import torch
import heapq

class HeapItem:
    def __init__(self, p, t):
        self.p = p
        self.t = t

    def __lt__(self, other):
        return self.p < other.p

h = []
heapq.heappush(h, HeapItem(1, torch.Tensor([[1,2]])))
heapq.heappush(h, HeapItem(1, torch.Tensor([[3,4]])))