Hi, I have two tensors tensor1 and tensor2, with shape (16, 768, 64, 64) and I want to stack them such that I have a tensor output with shape (16,2,768,64,64), BUT I also want the channel dimension 768 to be an alternating zip of tensor1 and tensor2.
tensor1[:, 0, :, :] and tensor2[:, 0, :, :] are stacked together as a (16, 2, 64, 64) tensor and so on and then we stack that list of stacked tensors again…
I don’t think there is a way to do this with torch.cat? So I wrote this little func that should do it for me
def zip_and_stack_tensors(tensor1, tensor2):
    num_of_channels = len(tensor1[0, :, 0, 0])
    imagelist = []
    for i in range(0, num_of_channels):
        tensorimage = torch.stack((tensor1[:, i, :, :], tensor2[:, i, :, :]), dim=1)
        imagelist.append(tensorimage)
    output = torch.stack(imagelist, dim=2)
return output

However, the backwardpass now takes about x8 times longer if I call this function; the forwardpass time does not change considerably. I guess because currently torch.stack gets called 768 + 1 times?
Is there a way to prevent this/write my function more efficient?