I have the following code:
input = torch.zeros(2,4,8448,2)
buf = torch.zeros(2,4,32,512,2)
blocks = torch.arange(32)
starts = blocks * 256
stops = starts + 512
for i in range(32):
buf[…,blocks[i], 0:512,:] = input[…,starts[i]:stops[i], :]
Now i would love to do:
buf[…,blocks,0:512,:] = input[…,starts:stops,:]
but this isn’t legal in pytorch.
Is there a cheat way of doing this? the for loop is a real killer