Мультикласс с NLLoss

При обучении автоэнкодера хочу использовать negative log likelihood loss, но при попытке обучения получаю ошибку о том, что loss функция не поддерживает multi-target. Пробовал NLLLoss2d, ошибка та же.

RuntimeError: 0D or 1D target tensor expected, multi-target not supported

Код обучения:

train_loader = DataLoader(train_ds, batch_size=256)

model = AEModel()
optimizer = Adam(model.parameters(), lr=1e-3)
loss = nn.NLLLoss2d(reduction="sum")

for epoch in range(EPOCHES):
    for data in train_loader:
        optimizer.zero_grad()

        output = model(data)

        loss = loss(output, data)
        loss.backward()

        optimizer.step()

Датасет:

np.random.choice([0, 1], size=(1024, 9724))