import torch
import torch.nn as nn
a1 = torch.randn((6,6,3,3))
a1[[1,2,3],[2,4,5],:,:] = torch.zeros((3,3))
print(a1)
这样只会有3个3*3被赋为0,一定要对应(1,2)(2,4) (3,5)
a1[range(3),range(3),:,:]也是 (0,0),(1,1)(2,2)
torch花式索引2
最新推荐文章于 2024-08-25 21:49:59 发布