torch中根据矩阵对应元素抽取矩阵中的值

下面是一个使用PyTorch的例子：

import torch


## 在这个函数中，我们首先将mask_tensor张量中值为1的位置转化为bool类型的掩码，然后使用这个掩码将data_tensor张量中对应位置的元素提取出来，并返回提取的元素。
def extract_elements_with_ones(mask_tensor, data_tensor):
    # 将mask_tensor张量中值为1的位置转化为bool类型掩码
    mask = mask_tensor == 1
    # 使用掩码将data_tensor张量中对应位置的元素提取出来
    extracted_data = data_tensor[mask]
    return extracted_data

# 示例输入
data_tensor = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
mask_tensor = torch.tensor([[0, 1, 0], [1, 0, 1], [0, 1, 0]])

# 调用函数提取元素
extracted_data = extract_elements_with_ones(mask_tensor, data_tensor)

# 打印提取的元素
print(extracted_data)
# tensor([2, 4, 6, 8])

在这个示例中，data_tensor是一个3x3的张量矩阵，mask_tensor是一个与data_tensor维度相同的张量矩阵，它的值为0或1。函数extract_elements_with_ones会提取data_tensor中与mask_tensor对应位置值为1的元素，最终返回一个包含所有提取元素的张量。