如何在两个不同大小的PyTorch张量中获得相等元素的索引？

import torch

# 创建两个张量
tensor_a = torch.tensor([[1, 2, 3], [4, 5, 6]])
tensor_b = torch.tensor([[4, 5, 6], [1, 2, 3]])

# 使用广播机制使两个张量形状相同
# 在这个例子中，tensor_a和tensor_b的形状已经是兼容的，所以不需要额外的广播步骤

# 比较两个张量并获取相等元素的索引
equal_indices = torch.nonzero(torch.eq(tensor_a, tensor_b), as_tuple=True)

print(equal_indices)

(tensor([0, 1, 0, 1, 0, 1]), tensor([1, 0, 2, 1, 2, 0]))

tensor_c = torch.tensor([[1], [2], [3]])
tensor_d = torch.tensor([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]])

# 将tensor_d的形状调整为(3, 4)
tensor_d_expanded = tensor_d.expand(-1, tensor_c.shape[1])

# 现在可以比较tensor_c和tensor_d_expanded
equal_indices_expanded = torch.nonzero(torch.eq(tensor_c, tensor_d_expanded), as_tuple=True)

print(equal_indices_expanded)