检测tensor中都为0的维度

你可以使用 PyTorch 中的 `torch.nonzero()` 函数来检测 tensor 中所有元素都为 0 的维度。

示例代码如下：

import torch

# 创建一个 2x3 的 tensor，其中第 2 行为全 0
tensor = torch.tensor([[1, 2, 3], [0, 0, 0]])

# 使用 nonzero 函数找到所有元素都为 0 的维度
zero_dims = tuple(torch.nonzero(tensor.sum(dim=1) == 0).squeeze().tolist())

print(zero_dims)  # 输出结果为 (1,)

在上述代码中，我们首先创建了一个 2x3 的 tensor，并将第 2 行全部置为 0。然后，我们使用 `tensor.sum(dim=1) == 0` 找到了所有元素都为 0 的行，再使用 `torch.nonzero()` 函数找到了这些行对应的维度。最后，我们将找到的维度转换为了一个元组，并打印出来。在这个例子中，因为只有第 2 行的所有元素都为 0，所以找到的维度只有一个，即 (1,)。