Given tensor a->[16,50,50,512] and b->[16,50,512]
I want this a[:,0,0,:]=b[:,0,:] from 0 to 49 with no for-loop
You should be able to use diagonal for assignment:
a = torch.zeros(6, 5, 5, 7)
b = torch.randn(6, 5, 7)
# the diagonal dimension is always last for numpy compat...
d = a.diagonal(dim1=1, dim2=2).permute(0, 2, 1)
d.copy_(b)
# now the "diagonal" of a is set to b.
Best regards
Thomas