As far as I understand you are dealing with a segmentation task, i.e. you would like to classify each pixel in your input image to one of your classes.
In your case you just have two valid classes (background and class).
If that’s correct, you could use a simple classification approach with your multi dimensional output and target.
Here is a dummy example using random input and a random target consisting of 2 classes:
batch_size = 2
c, h, w = 3, 24, 24
nb_classes = 2
x = torch.randn(batch_size, c, h, w)
target = torch.empty(batch_size, h, w, dtype=torch.long).random_(nb_classes)

model = nn.Sequential(
    nn.Conv2d(c, c*2, 3, 1, 1),
    nn.ReLU(),
    nn.Conv2d(2*c, nb_classes, 3, 1, 1)
)
criterion = nn.CrossEntropyLoss()

output = model(x)
loss = criterion(output, target)
loss.backward()