Hi,
I’m not sure to understand what you’re trying to do here. You can get the same result as your zip_and_stack_tensors() using a single torch.stack(). Am I missing something?
import torch
from torch import nn, optim

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

t1 = torch.rand(16, 768, 64, 64)
t2 = torch.rand(16, 768, 64, 64)

out = zip_and_stack_tensors(t1, t2)
out_opt = torch.stack([t1, t2], 1)

print(out.size() == out_opt.size())
print((out - out_opt).abs().max())