import torch
import os
import numpy as np
import torch.nn as nn
import torch.optim as optim

# DataLoader
from PIL import Image
from torchvision import transforms
from torch.utils.data.sampler import SubsetRandomSampler

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()

        self.relu = nn.ReLU()
        self.conv1 = nn.Conv2d(3, 64, (5, 5), (1, 1), (2, 2))
        self.conv2 = nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1))
        self.conv3 = nn.Conv2d(64, 32, (3, 3), (1, 1), (1, 1))
        self.conv4 = nn.Conv2d(32, 1, (3, 3), (1, 1), (1, 1))

    def forward(self, x):
        x = self.relu(self.conv1(x))
        x = self.relu(self.conv2(x))
        x = self.relu(self.conv3(x))
        x = self.conv4(x)
        return x


img  = torch.randn(1,3, 250, 250, requires_grad = True)
mask  = torch.randn(1,1, 250, 250, requires_grad = True)
# Model:
from torchvision import models
model = Net()
criterion = nn.BCEWithLogitsLoss()
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
model.train()


# Use gpu if available
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
if torch.cuda.device_count() > 1:
    print("Let's use", torch.cuda.device_count(), "GPUs!")
    model = nn.DataParallel(model)

model.to(device)


# Training
num_epochs = 1
for iter in range(num_epochs):
    # Training:
    model.train() 
    for i in range(10):
        # img, mask = train_loader[i]
        img, mask = img.cuda(), mask.cuda()
        
        optimizer.zero_grad()
        img = img.to(device)
        mask = mask.to(device)
        output = model(img)
        loss = criterion(output, mask)
        print(loss.cpu())
        loss.backward()
        optimizer.step()
    
    # Validation:
    model.eval() 
    for i in range(10):
        img, mask = test_loader[i]
        img, mask = img.cuda(), mask.cuda()
        output = model.forward(img) # forward pass
        loss = criterion(output, mask)
        print(loss.cpu())


This is a toy code based on yours. Hope this helps