Implementing CrossEntropyLoss

Here is a simple model I wrote below that utilizes CrossEntropyLoss :

f = []
f.append([1,1,1])
f.append([0,1,0])
f.append([0,1,0])
x = torch.FloatTensor(np.array(f))

t = []
t.append(0)
t.append(1)
t.append(1)
y = torch.LongTensor(np.array(t))

dimensions_input = 3
hidden_layer_nodes = 3
output_dimension = 3

class Model(torch.nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.linear = torch.nn.Linear(dimensions_input,hidden_layer_nodes)
        self.linear2 = torch.nn.Linear(hidden_layer_nodes,output_dimension)

    def forward(self, x):
        l_out1 = self.linear(x)
        y_pred = self.linear2(l_out1)
        return y_pred

model = Model()

criterion = torch.nn.CrossEntropyLoss()
optim = torch.optim.SGD(model.parameters(), lr = 0.0001)

def main():
    for i in range(50000):
        y_pred = model(x)
        loss = criterion(y_pred, y)

        if i % 1000 == 0:
            print(loss)
        optim.zero_grad()

        loss.backward()
        optim.step()
        
main()

I achieve a low loss value for this implementation : tensor(0.1186)

But when i attempt to predict a value using model(x[0]) the returned value is not close to ground truth of [1,1,1] as value : tensor([ 2.7715, 0.7934, -1.1497]) is returned.

Have I implemented this model correctly ?