如何在PyTorch张量的if语句中设置多条件并逐行判断

你的问题在于原代码直接对整个张量的布尔结果做判断，无法实现逐行检查的需求。要逐行判断每行是否满足第一列<1且第二列<1的条件，可以按以下方式处理：

实现代码 #

import torch
all_data = torch.tensor([[1.1, 0.4],
                         [1.7, 2.7],
                         [0.9, 0.7],
                         [0.9, 3.5],
                         [0.1, 0.5]])

# 生成逐行的布尔掩码，标记每行是否满足条件
row_mask = (all_data[:, 0] < 1) & (all_data[:, 1] < 1)

# 逐行遍历并输出结果
for row_idx, is_match in enumerate(row_mask):
    if is_match:
        print(f"第{row_idx+1}行: Masked")
    else:
        print(f"第{row_idx+1}行: Not_Masked")

代码说明 #

1. 生成掩码：通过(all_data[:, 0] < 1) & (all_data[:, 1] < 1)得到一个布尔张量row_mask，其中每个元素对应原张量的一行，True表示该行满足条件，False则不满足。
2. 逐行判断：用enumerate遍历掩码，根据每个位置的布尔值输出对应的结果。

如果需要提取满足条件的张量行，可以直接使用all_data[row_mask]获取这些数据。

内容的提问来源于stack exchange，提问作者Link_tester