This is correct! Below is an example of the collate_fn I was using to batch together examples from the Dataset:
def myCollater(examples):
    input_imgs = []
    target_dicts = []
    for ex in examples:
        input_imgs.append(ex[0])
        target_dicts.append(ex[1])
        
    return torch.stack(input_imgs), target_dicts

I ended up appending the dictionaries to a list that came directly from the Dataset. The input image tensors didn’t stay on the GPU when training because I thought that torch.stack() basically combines multiple image tensors into one single batched tensor (which is a copy of all of the original tensors, so putting it on the GPU isn’t necessarily putting the original image tensors on the GPU).
Also, if we wanted to train a Mask_RCNN with variable image sizes, then the collate function should return a list of image tensors (rather than a single stacked/batched image tensor, since you can’t stack multiple tensors of variable size). This meant I had to switch my return function of myCollater to:
return input_imgs, target_dicts
This also means that I have to now put the images on the GPU one at a time (since input_imgs is a list of image tensors), so in the training loop we now have:
for j in range(len(images)):
    images[j] = images[j].to(device)

instead of:
images = images.to(device)

I thought that this would now cause the images to stay on the GPU, since they are supplied directly from the Dataset class, and appended in a list through the myCollater function. Turns out that these images don’t stay on the GPU, and everything runs fine.
I ran a small experiment to find out that when you add a tensor to a list, a copy of the tensor is made every time it’s added to a list. (So each list has it’s own unique reference to a tensor). For dictionary values, it’s the same thing!

image1323×752 125 KB

So my GPU memory problem was caused from appending original dictionaries of tensors. Appending dictionaries (or lists) doesn’t necessarily copy the objects for you automatically, which is why we have to do it ourselves explicitly in the training loop (or in the collate_fn). Oops!
Thanks for the info again