Here is the code I have:
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import transforms
from torch.utils.data import DataLoader, Dataset, TensorDataset

bs=1
from torchvision.models import resnet18
model = resnet18(pretrained=False)
optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)

input = torch.rand(bs,3, 256,256)
target = torch.rand(bs, 1000)
print(target.size())

# target = target.long()
model.train()

# before we learn
p1 = model.fc.weight
od1 = model.state_dict()


output = model(input)
# print(output.size())
loss_fn = torch.nn.CrossEntropyLoss()
# loss_fn = torch.nn.BCEWithLogitsLoss()
loss = loss_fn(output, target)
print(loss)
if optimizer is not None:
        optimizer.zero_grad()
        # backward pass + optimize
        loss.backward()
        optimizer.step()

p2 = model.fc.weight
od2 = model.state_dict()

print(torch.equal(p1,p2))

torch.nn.CrossEntropyLoss() should be the main loss function (criterion) for ImageNet classification. I just tried to test the code and got this error I could not fix fast.
What I am expecting at the end is that print(torch.equal(p1,p2)) will return False. This means that the parameter before and after the single optimization step will be different ( the model learned something ).
If I set the torch.nn.BCEWithLogitsLoss()  it will work but the params p1 and p2 will be just the equal.
What am I doing wrong?