Hi @ptrblck, I tried working with a more even dataset with balanced classes. I also tried playing around with the weight_decay parameter in the optimizer. However, I am still running into the issue where my model is only outputting 1 class. I finally decided to just increase the epochs and overfit the model with training data and perform prediction on a training image. It turns out I am still only getting 1 class when I am supposed to get 2 classes. I tried scouring the internet for a similar issue, but couldn’t find one.  Any and all help is very much appreciated.
Thanks
Nishanth
I have provided a detailed explanation in the following post

  
    
    
    Semantic Segmentation using deeplabv3+resnet101 from torchvision models vision
  
  
    I am using the Deeplab V3+ resnet 101 to perform binary semantic segmentation. 
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.9
w_decay    = 1e-5
step_size  = 50
gamma      = 0.5
traindata = …