python - 使用重复的优先级将 Torch 张量插入 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]])))