PyTorch - Replace a log_softmax funtion by F.cross_entropy with different result

I'm working on the PyTorch tutorial "What is torch.nn really?". There a simple MNIST neural net is developed from scratch and then gets refactored piece for a piece with torch elements. It is used a log_softmax function:

def log_softmax(x):
    return x - x.exp().sum(-1).log().unsqueeze(-1)

def model(xb):
    return log_softmax(xb @ weights + bias)

After finishing the simple neural net the first replacement is above function by (in a continuous Jupyter cell):

import torch.nn.functional as F
loss_func = F.cross_entropy

def model(xb):
    return xb @ weights + bias

The proof of the result is fine with a further call of

print(loss_func(model(xb), yb), accuracy(model(xb), yb))

which brings in both cases the same result, e. g.

loss_func 1: tensor(0.0830, grad_fn=) tensor(1.)
vs. loss_func 2: tensor(0.0830, grad_fn=) tensor(1.)

Now I tried to make the replacement directly there where the definition of the log_softmax function is done, expecting a comparable result with the proof functions. But the result is totally different

loss_func 1: tensor(-3578.4377, grad_fn=) tensor(0.5625) vs.
loss_func 2: tensor(310.3191, grad_fn=) tensor(0.5625)

Can you please help me find my mistake?

Here is the complete code. It's close to the original, only the the comments are removed. The log_softmax func is outcommented and replaced by the F.cross_entropy.

from pathlib import Path
import requests

DATA_PATH = Path("data")
PATH = DATA_PATH / "mnist"

PATH.mkdir(parents=True, exist_ok=True)

URL = "https://github.com/pytorch/tutorials/raw/master/_static/"
FILENAME = "mnist.pkl.gz"

if not (PATH / FILENAME).exists():
        content = requests.get(URL + FILENAME).content
        (PATH / FILENAME).open("wb").write(content)

import pickle
import gzip

with gzip.open((PATH / FILENAME).as_posix(), "rb") as f:
        ((x_train, y_train), (x_valid, y_valid), _) = pickle.load(f, encoding="latin-1")

from matplotlib import pyplot
import numpy as np

pyplot.imshow(x_train[0].reshape((28, 28)), cmap="gray")
print(x_train.shape)

import torch

x_train, y_train, x_valid, y_valid = map(
    torch.tensor, (x_train, y_train, x_valid, y_valid)
)
n, c = x_train.shape
print(x_train, y_train)
print(x_train.shape)
print(y_train.min(), y_train.max())

import math

weights = torch.randn(784, 10) / math.sqrt(784)
weights.requires_grad_()
bias = torch.zeros(10, requires_grad=True)

"""
def log_softmax(x):
    return x - x.exp().sum(-1).log().unsqueeze(-1)

def model(xb):
    return log_softmax(xb @ weights + bias)
"""

import torch.nn.functional as F
loss_func = F.cross_entropy

def model(xb):
    return xb @ weights + bias

bs = 64  # batch size

xb = x_train[0:bs]  # a mini-batch from x
preds = model(xb)  # predictions
preds[0], preds.shape
print(preds[0], preds.shape)

def nll(input, target):
    return -input[range(target.shape[0]), target].mean()

loss_func = nll

yb = y_train[0:bs]
print(loss_func(preds, yb))

def accuracy(out, yb):
    preds = torch.argmax(out, dim=1)
    return (preds == yb).float().mean()

print(accuracy(preds, yb))

from IPython.core.debugger import set_trace

lr = 0.5  # learning rate
epochs = 2  # how many epochs to train for

for epoch in range(epochs):
    for i in range((n - 1) // bs + 1):
        #         set_trace()
        start_i = i * bs
        end_i = start_i + bs
        xb = x_train[start_i:end_i]
        yb = y_train[start_i:end_i]
        pred = model(xb)
        loss = loss_func(pred, yb)

        loss.backward()
        with torch.no_grad():
            weights -= weights.grad * lr
            bias -= bias.grad * lr
            weights.grad.zero_()
            bias.grad.zero_()

print("loss_func 1: ", loss_func(model(xb), yb), accuracy(model(xb), yb))

import torch.nn.functional as F
loss_func = F.cross_entropy

def model(xb):
    return xb @ weights + bias

print("loss_func 2: ", loss_func(model(xb), yb), accuracy(model(xb), yb))