Python Pytork将操作应用于某些行，但不是全部行

我试图在张量的某些行上应用softmax函数，但问题是我的某些行具有all-inf值。因此，这些行上的softmax会输出NaN，这会导致模型中出现问题

---

因此，我想创建一个函数，将softmax应用于一行，除非in是all-inf。在这种情况下，它输出一个零向量。有什么简单的方法可以做到这一点吗？

是否需要将softmax之后所有

nan

的行设置为0？这样可以确保不会覆盖任何意外的NAN

import numpy as np
import torch
import torch.nn.functional as F

array = np.arange(25, dtype=np.float32).reshape((5, 5))
array[3, ...] *= -np.inf

# [[  0.   1.   2.   3.   4.]
#  [  5.   6.   7.   8.   9.]
#  [ 10.  11.  12.  13.  14.]
#  [-inf -inf -inf -inf -inf]
#  [ 20.  21.  22.  23.  24.]]

array = torch.tensor(array)
array = F.softmax(array, dim=1)
mask = array.isnan().all(dim=1)
array[mask, ...] = 0
print(array)

# tensor([[0.0117, 0.0317, 0.0861, 0.2341, 0.6364],
#         [0.0117, 0.0317, 0.0861, 0.2341, 0.6364],
#         [0.0117, 0.0317, 0.0861, 0.2341, 0.6364],
#         [0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
#         [0.0117, 0.0317, 0.0861, 0.2341, 0.6364]])

---

在做这件事之前，你应该问问自己，为什么你首先会得到

-inf

。如果这是训练的结果，你在其他地方有问题需要解决。如果它们已经存在，您应该在培训之前将它们从数据集中删除-这是标准做法。