Hi! Now I want to pass a 5-d tensor to conv layers, where the the tensor is of [k, batch_size, C,H,W]. The later 4 dimension [batch_size,C,H,W] is just the usual 4-D tensor of image input, but this 4-D tensor is copied and expanded k times along the first dimension. Such an input is necessary in models like IWAE. I want to do usual 2D convolution on this tensor, while making sure that the forward pass output for each of the k copies are the same. I think reshaping it to [k*batch_size,C,H,W] will make the forward output for each copy different, so is there any smart way to implement this?

I think reshaping it to `k*batch_size,C,H,W`

would work. you can check this out with for example `batch_size =1`

, and `k=3`

, the output would be `3,C,H,W`

then, and you can verify that `0,:,:,:`

is the same as `1,:,:,:`

1 Like