Linear classifier with PyTorch from scratch questions

I'm trying to implement linear classifier in PyTorch, using 1 layer with tensors W and b, softmax and cross entropy loss. For each batch I have to:

1. Calculate logits
2. Transform logits to probabilities with softmax
3. Compute most probable classes
4. Calculate cross entropy between true and predicted classes
5. Use an optimizer to change W and b

So far I have (I have flat MNIST loaded with Scikit-learn):

# convert Numpy arrays to PyTorch tensor Variables
input_X_train = torch.from_numpy(X_train_flat).float().to(device)
input_X_val = torch.from_numpy(X_val_flat).float().to(device)
input_X_test = torch.from_numpy(X_test_flat).float().to(device)

input_y_train = torch.from_numpy(y_train).long().to(device)
input_y_val = torch.from_numpy(y_val).long().to(device)
input_y_test = torch.from_numpy(y_test).long().to(device)

# model parameters: W and b
W = torch.randn(input_dim, output_dim, device=device, dtype=dtype, requires_grad=True)
b = torch.randn(1, device=device, dtype=dtype, requires_grad=True)

BATCH_SIZE = 512
EPOCHS = 40
LEARNING_RATE = 1e-6

# create torch.optim.Adam optimizer for loss function minimization
optimizer = torch.optim.Adam([W, b], lr=LEARNING_RATE)

# create negative log loss function object for loss function evaluation
# use mean loss value from all batch samples
loss_fn = torch.nn.NLLLoss(reduction="mean")

for t in range(EPOCHS):    
    # logits for input_X, resulting shape should be [input_X.shape[0], 10]
    logits = torch.matmul(input_X_train, W) + b

    # apply torch.nn.functional.softmax (torch_F.softmax) to logits
    probas = torch_f.softmax(logits, dim=1)
    
    # apply torch.argmax to find a class index with highest probability
    classes = torch.argmax(probas, dim=1)

    # loss should be a scalar number: average loss over all the objects with torch.mean()
    # PyTorch implements negative log loss (NLL) *without* log - you have to first compute log of 
    # softmax, then negative log loss, which will swap sign
    
    # Use torch.nn.functional.log_softmax (torch_f.log_softmax) on top of input_y and logits
    # It is identical to calculating cross-entropy (log and then NLL) on top of probas, 
    # but is more numerically friendly (read the docs).
    log_probas = torch_f.log_softmax(logits, dim=1)
    loss = loss_fn(log_probas, input_y_train)

    # Before the backward pass, use the optimizer object to zero all of the
    # gradients for the variables it will update (which are the learnable
    # weights of the model). This is because by default, gradients are
    # accumulated in buffers( i.e, not overwritten) whenever .backward()
    # is called. Checkout docs of torch.autograd.backward for more details.
    optimizer.zero_grad()
    
    # calculate backward gradients for backpropagation
    loss.backward()
    
    # Calling the step function on an Optimizer makes an update to its parameters
    optimizer.step()

For some reason, the W and b don't change. What am I doing wrong?

EDIT: I've seen and tried in the code above e. g. this minimal working example https://discuss.pytorch.org/t/minimal-working-example-of-optim-sgd/11623/2.

EDIT 2: Gradients W.grad are often, I think it should not be like that. Probabilities of classes are definitely right (so it's not e. g. like this example), since I've checked sum of every row and probabilities of all classes for each sample sum to 1.