将张量中的inf值替换为指定数据

将张量中的inf值替换为指定数据

代码

import torch

x1 = [[[[1, 2], [3, 0]]]]
x2 = [[[[1, 2], [3, 4]]]]
x1 = torch.tensor(x1)
x2 = torch.tensor(x2)
x3 = x2/x1
print(x3)

x3_output = torch.where(x3 == float('inf'), 382, x3)
# x3_output = torch.where(x3 == float('-inf'), 382, x3)
print(x3_output)

运行结果