I’m trying to build a logistic regression model using pytorch for training a learning to rank model using a custom listwise loss function. The issue is the the loss isn’t changing in the 10 epochs and the softmax output (variable output) comes out to be 1 for all rows. Please find the code below.
STEP 1: CREATE MODEL CLASS
class LogisticRegressionModel(nn.Module):
     def __init__(self, input_size, num_classes):
         super(LogisticRegressionModel, self).__init__()
         self.linear = nn.Linear(input_dim, output_dim)

    def forward(self, x):
        out = F.softmax(self.linear(x))
        return out

STEP 2: INSTANTIATE MODEL CLASS
input_dim = X_train.shape[1]
output_dim = 1
n_unique_qid = 119
n=list_size
model = LogisticRegressionModel(input_dim, output_dim)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model.to(device)

#custom loss function
def listwiseloss(output):
    s=0
    b=0
    for i in range(n_unique_qid):
      for j in range(list_size[i]-1):
        s=s-output[sum(list_size[:i])+j]
        for k in range(j,list_size[i]):
          b =b+math.exp(output[sum(list_size[:i])+k])
        s = s+math.log(b)
        b=0  
    return s


STEP 4: INSTANTIATE OPTIMIZER CLASS
from torch.autograd import Variable


learning_rate = 0.01

optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)


STEP 5: TRAIN THE MODEL
epochs = 10
for epoch in range(epochs):
    epoch += 1
    # Convert numpy array to torch Variable

    #######################
    #  USE GPU FOR MODEL  #
    #######################
    train = torch.tensor(X_train, dtype=torch.float, device=device)
    #Variable(torch.ones(2, 2), requires_grad=True)
    #train = torch.from_numpy(X_train).to(device)
    inputs = Variable(train, requires_grad=True)
    #labels = torch.from_numpy(y_train).to(device)

    # Clear gradients w.r.t. parameters
    optimizer.zero_grad() 

    # Forward to get output
    outputs = model(inputs)

    # Calculate Loss
    loss = listwiseloss(outputs)

    # Getting gradients w.r.t. parameters
    loss.backward()

    # Updating parameters
    optimizer.step()

    # Logging
    print('epoch {}, loss {}'.format(epoch, loss.item()))

What could be the reason behind it? Really in need of an expert advice since this is an internship project and needs to be soon.