是否有一个 torch 函数可以导出两个张量的并集？

Question

numpy中有一个函数可用于导出两个张量的并集，如下所示：

import torch
import numpy as np
a = torch.tensor([0, 1, 2])
b = torch.tensor([2, 3, 4])
c = np.union1d(a, b) # c = array([0, 1, 2, 3, 4])
c = torch.from_numpy(c) # c = torch.tensor([0, 1, 2, 3, 4])

但是，我正在寻找可以直接在两个张量上使用的torch函数。如果我使用上面的 numpy 函数，我必须将结果从 numpy 转换为 torch，并且还必须使用 cpu，尽管该函数应用于张量。

torch中有没有可以直接在两个张量上使用的联合函数？或者，至少可以使用其他火炬功能简单地实现它吗？

Answer 1

您可以使用：

torch.cat((a, b)).unique()