I am very new to Pytorch and Deep Learning in general. I have a set of grayscale images which I convert to a 3 channel images by repeating the first channel two more times. I am using the fcn_resnet101 from the pytorch models. The model only predicts one class for all images. Am I missing something from the methodology below.
import torch
import torchvision
import loader
from loader import DataLoaderSegmentation
import torch.nn as nn
import torch.optim as optim
import numpy as np
from torch.utils.data.sampler import SubsetRandomSampler
batch_size = 1
validation_split = .2
shuffle_dataset = True
random_seed= 66

n_class    = 2
num_epochs = 1
#lr         = 1e-4
#momentum   = 0
w_decay    = 1e-5
step_size  = 50
gamma      = 0.5

traindata = DataLoaderSegmentation('/home/ubuntu/Downloads/Brain/test0716')
dataset_size = len(traindata)
indices = list(range(dataset_size))
split = int(np.floor(validation_split * dataset_size))
#if shuffle_dataset:
np.random.seed(random_seed)
np.random.shuffle(indices)
train_indices = indices[split:]
val_indices = indices[:split]

#train_sampler = SubsetRandomSampler(train_indices)
#valid_sampler = SubsetRandomSampler(val_indices)
trainloader = torch.utils.data.DataLoader(traindata, batch_size=batch_size)
validation_loader = torch.utils.data.DataLoader(traindata, batch_size=batch_size)

#trainloader = torch.utils.data.DataLoader(traindata, batch_size=1, shuffle=False, num_workers=2)

model = torchvision.models.segmentation.fcn_resnet101(pretrained=False, progress=True,aux_loss=None, num_classes=2).cuda()
criterion = nn.CrossEntropyLoss().cuda()
#optimizer = optim.RMSprop(model.parameters(), lr=lr, momentum=momentum, weight_decay=w_decay)
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.train()


for iter in range(num_epochs):
    print(iter)	
    for (i,l) in trainloader:
        i= i.to(device)
        l = l.to(device=device, dtype=torch.int64)
        outt = model(i)
        loss = criterion(outt['out'], l.squeeze(1))
        #print(loss)
        loss.backward()
        optimizer.step()
torch.save(model, '/home/ubuntu/Downloads/Brain/test0716/mod.pth)

Thanks for all your help