I am using a softmax layer at the end of my fully connected network. It does not sum to 1. This is an issue because it tends to increase proportional to the number of epochs. Not sure why this is happening here is the code.
class NN(nn.Module):
    def __init__(self):
        super(NN, self).__init__()

        self.fc1 = nn.Linear(1024, 128)
        self.fc2 = nn.Linear(128, 128)
        self.softmax = nn.Softmax(0)

    def forward(self, input):
        output = self.fc1(input)
        output = self.fc2(output)
        ouput = self.softmax(output)
        return output



def train(x_tensor, y_tensor):
    output = rnn(x_tensor.view(-1))
    output = output.unsqueeze(0)
    y_tensor = y_tensor[0].unsqueeze(0)
    y_guess = np.argmax(np.array(y_tensor))  #.long()
    y_guess = torch.from_numpy(np.array(y_guess)).unsqueeze(0)
    loss = criterion(output, y_guess)  
    loss.backward()
    optimizer.step()
    return output, loss.item()


      rnn = NN()
      lr = 0.001
      optimizer = torch.optim.SGD(rnn.parameters(), lr = lr)
      criterion = nn.CrossEntropyLoss()