If the original image resolution yields an out of memory error on your device, you could either try to:

apply torch.utils.checkpoint to trade compute for memory
use model sharding, if you have more than a single device, and execute parts of the model on separate devices
execute the last iteration on the CPU, since you would usually have more system RAM than GPU memory
try mixed precision training 1 for a potentially reduced memory usage and speedup (you would need to install the nightly binaries for this or build from master).