pytorch张量** 2计算

Question

for i in range(180000):
    p[indices[0][i], :, indices[1][i], indices[2][i]] = \
        p[indices[0][i], :, indices[1][i], indices[2][i]] ** 2

需要60秒钟以上的时间，p = p ** 2可以立即获得结果。我只想更改 P 的一部分。那么，我应该怎么做才能达到这种效率呢？通过 Pytorch 编码。

Answer 1

也许我已经知道它的原因。我应该将索引转换为cuda格式。