Large image for semantic segmentation task — Pytorch runs out of memory

Question

I have large images (256,256) of ImageNet, and I want to train for a semantic segmentation task. The output will thus be (B,Classes,C,H,W). I have 256 classes. For