Actually what I got was the following error, sorry for not being specific:
output += segment
RuntimeError: output with shape [1, 3, 1000, 1000] doesn't match the broadcast shape [8, 3, 1000, 1000]

I was able to fix this issue by changing
output = torch.zeros(1, 3, out.size(-2), out.size(-1), dtype=torch.float)

to
output = torch.zeros(out.shape[0], 3, ... ) #this ensures that regardless of the batch size, output image shape will match up

But now my output image consists of a grid of size batch (8)… So output contains 8 images, do you know how I can iterate through the first dimension and get each of the 8 images from it? Thank you very much!