如何移除PyTorch张量中特定维度含NaN的所有索引？

需求明确：移除张量中所有在同一行（第二维度）存在NaN的行——即只要某一行在任意样本中包含NaN，就删除所有样本的该行。

实现步骤及代码 #

1. 生成NaN掩码，检测每一行是否存在NaN：

import torch

# 补全原张量的语法错误后定义
a = torch.tensor([[[2,3],
                   [3,4],
                   [4,5]],
                  [[3,6],
                   [6,2],
                   [2,-1]],
                  [[float('nan'), 1],
                   [2,3], 
                   [3,2]]])

# 沿第一、第三维度检查，标记第二维度中存在NaN的行
has_nan = torch.isnan(a).any(dim=(0, 2))

2. 筛选保留行并生成干净张量：

# 取反得到需要保留的行索引
keep_rows = ~has_nan

# 通过索引得到最终结果
a_clean = a[:, keep_rows, :]

验证结果 #

执行后a_clean的输出与预期完全一致：

tensor([[[3, 4],
         [4, 5]],

        [[6, 2],
         [2, -1]],

        [[2, 3],
         [3, 2]]])

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