Ошибка времени выполнения: ожидался скалярный тип Long, но найден Float

Описание проблемы:

Я сталкиваюсь с проблемой несовпадения типов данных между выходом модели и метками при вычислении функции потерь. В зависимости от того, в какой тип я изменяю тензоры (долгое или плавающее), возникают ошибки.

Имеющиеся тензоры имеют форму (42000, 1, 28, 28) для изображений и (42000) для меток. Я не уверен, где можно изменить необходимые типы данных для модели или функции потерь.

Кроме того, я не уверен, нужен ли мне DataLoader, так как использование Variable тоже не дало результата.

Вот мой код:

dataloaders_train = torch.utils.data.DataLoader(Xt_train, batch_size=64)

dataloaders_test = torch.utils.data.DataLoader(Yt_train, batch_size=64)

class Network(nn.Module):
    def __init__(self):
        super().__init__()
        self.hidden = nn.Linear(42000, 256)
        self.output = nn.Linear(256, 10)
        self.sigmoid = nn.Sigmoid()
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x):
        x = self.hidden(x)
        x = self.sigmoid(x)
        x = self.output(x)
        x = self.softmax(x)
        return x

model = Network()

input_size = 784
hidden_sizes = [28, 64]
output_size = 10 
model = nn.Sequential(nn.Linear(input_size, hidden_sizes[0]),
                      nn.ReLU(),
                      nn.Linear(hidden_sizes[0], hidden_sizes[1]),
                      nn.ReLU(),
                      nn.Linear(hidden_sizes[1], output_size),
                      nn.Softmax(dim=1))

criterion = nn.NLLLoss()
optimizer = optim.SGD(model.parameters(), lr=0.003)

epochs = 5

for e in range(epochs):
    running_loss = 0
    for images, labels in zip(dataloaders_train, dataloaders_test):
        images = images.view(images.shape[0], -1)
        print(images.dtype)
        print(labels.dtype)

        optimizer.zero_grad()
        output = model(images)
        loss = criterion(output, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()
    else:
        print(f"Training loss: {running_loss}")

При выполнении кода я получаю следующую ошибку:

RuntimeError: expected scalar type Long but found Float

Как мне решить данную проблему? Как можно убедиться, что типы данных совпадают, чтобы избежать этой ошибки?