Yes, the snippet does reconstruct the patched image, albeit one patch at a time. So you’re suggesting something like the following for batch processing:
reshaped = patches.view(-1,k,k) #(batch,h,w)
temp = torch.empty(reshaped.shape)
loader = Dataloader(dataset=reshaped,batch_size=16)
for batch in loader:
   x = batch.view(-1,1,k,k).to(device, dtype = torch.float)
   temp[batch_idx,:,:] = model(x)

I don’t understand how to get the patch index from each batch to fill the empty tensor (temp). Can you clarify how that should work? Thanks!