I am trying to get a simple network to output the probability that a number is in one of three classes. These are, smaller than 1.1, between 1.1 and 1.5 and bigger than 1.5. I am using cross entropy loss with class labels of 0, 1 and 2, but cannot solve the problem.
Every time I train, the network outputs the maximum probability for class 2, regardless of input. The lowest loss I seem to be able to achieve is 0.9ish. Any advice on where I am going wrong would be greatly appreciated!! All code is below.
class gating_net(nn.Module):
    def __init__(self, input_dim, output_dim):
        super(gating_net, self).__init__()
        self.linear1 = nn.Linear(input_dim, 32)
        self.linear2 = nn.Linear(32, output_dim)

    def forward(self, x):
        # The original input (action) is used as the residual.
        x = F.relu(self.linear1(x))
        x = F.sigmoid(self.linear2(x))
        return x

learning_rate = 0.01
batch_size = 64
epochs = 500
test = 1

gating_network = gating_net(1,3)

optimizer = torch.optim.SGD(gating_network.parameters(), lr=learning_rate, momentum=0.9)
scheduler = ReduceLROnPlateau(optimizer, mode='max', factor=0.5, patience=20, verbose=True)

for epoch in range (epochs):
    input_ = []
    label_ = []

    for i in range (batch_size):
        scale = random.randint(10,20)/10

        input = scale
        if scale < 1.1:
            label = np.array([0])
        elif 1.1 < scale < 1.5:
            label = np.array([1])
        else:
            label = np.array([2])

        input_.append(np.array([input]))
        label_.append(label)

    optimizer.zero_grad()

    # get output from the model, given the inputs
    output = gating_network.forward(torch.FloatTensor(input_))
    old_label  = torch.FloatTensor(label_)

    # get loss for the predicted output
    loss = nn.CrossEntropyLoss()(output, old_label.squeeze().long())

    # get gradients w.r.t to parameters
    loss.backward()
    # update parameters
    optimizer.step()
    scheduler.step(loss)

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

    if loss.item() < 0.01:
        print("########## Solved! ##########")
        torch.save(mod_network.state_dict(), './supervised_learning/run_{}.pth'.format(test))
        break

    # save every 500 episodes
    if epoch % 100 == 0:
        torch.save(gating_network.state_dict(), './run_{}.pth'.format(test))