I am implementing a non-linear regression using neural networks with one single layer in Pytorch. However, using an activation function as ReLu or Softmax, the loss gets stuck, the value does not decrease as the sample increases and the prediction is constant values. So, I replaced ReLu, with LeakyReLU, and the loss decreased substantially, and the predictions were no longer constant and even tracked the original function.

However, in the context, which I am working, the Softmax function would be more appropriate. However, the vanishing gradient problem persists. I have tried to initialize with small weights but it does not work. I am wondering if someone could give me an idea on how to increase the steepness of the Softmax function on Pytorch since it worked with LeakyReLU.

class NeuralNetwork(nn.Module):
    def __init__(self,inputsize,outputsize):
        super(NeuralNetwork, self).__init__()
        self.flatten = nn.Flatten()
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(inputsize, outputsize),


    def forward(self, x):
        x = self.flatten(x)
        logits = self.linear_relu_stack(x)
        return logits

The hyperparameters that I am using are the following:

inputDim = 1        # takes variable 'x' 
outputDim = 1       # takes variable 'y'
learningRate = 0.001 
epochs = 100000
model = NeuralNetwork(inputDim, outputDim)
if torch.cuda.is_available():

If it is needed I can provide the simulated data.

criterion = torch.nn.MSELoss() 
optimizer = torch.optim.SGD(model.parameters(), lr=learningRate)
for epoch in range(epochs):
    # Converting inputs and labels to Variable
    if torch.cuda.is_available():
        inputs = Variable(torch.from_numpy(vS0).cuda().float())
        labels = Variable(torch.from_numpy(vC).cuda().float())
        inputs = Variable(torch.from_numpy(vS0).float())
        labels = Variable(torch.from_numpy(vC).float())


    # get output from the model, given the inputs
    outputs = model(inputs)

    # get loss for the predicted output
    loss = criterion(outputs, labels)

    # get gradients w.r.t to parameters

    # update parameters

    print('epoch {}, loss {}'.format(epoch, loss.item()))
  • $\begingroup$ This does not look like a vanishing gradient problem. Maybe it has to do with the loss function. Can you show us that part of the code? $\endgroup$
    – noe
    Sep 27, 2022 at 16:03
  • $\begingroup$ @noe I have updated the remaining code, except for the simulation of the data. I am using the mean square distance as a loss function. $\endgroup$ Sep 27, 2022 at 16:45
  • $\begingroup$ Softmax is meant normally for multiclass classification problems. Why do you want to use it for regression? $\endgroup$
    – noe
    Sep 29, 2022 at 14:42
  • $\begingroup$ @noe I wanted a smooth function since the output produces a smooth graph. The softmax was advised from a purely theoretical perspective. $\endgroup$ Oct 1, 2022 at 13:12

1 Answer 1


Softmax is the multiclass version of Cross-Entropy Loss, which is known to have vanishing gradient problems. Change the Softmax activation to ReLU.

  • $\begingroup$ It only works with LeakyReLU. I wanted a smooth function since the output produces a smooth graph. The softmax was advised from a purely theoretical perspective. $\endgroup$ Oct 1, 2022 at 13:11

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge that you have read and understand our privacy policy and code of conduct.

Not the answer you're looking for? Browse other questions tagged or ask your own question.