This is exactly the code that I have now, but I just realized I had passed another model in the training function 
I got a few errors I could fix, but now I’m stuck on this one that I moved to a new topic: Target size (torch.Size([3, 3, 256, 256])) must be the same as input size (torch.Size([3, 65536]))
Thank you so much for the help though!