Hi ptrblck, thank you so much for your elaborate answer.
When I started, I declared the number of classes as 1 when initializing my deeplab model. The issue was that the values in the output ranged between [-1.5, 1.5]. And taking the argmax of this occurred in only zeros, nothing I could really use for my accuracy measurements.
Defining as I did at start with 1 class seems to be the easier method as it matches the shape with the mask, but I’m not sure how I’ll  find the threshold if my values are in a continuous range of [-1, 1.5].
If you don’t mind, could you give me some guidance over why my model output would even be negative and above 1?
Maybe the issue is at my custom dataset. First I was just applying ToTensor() which normalized values between [0, 1] (this still produces [-1, 1.5] range from model output, however, after applying transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])]), they started to range to negative values even before training the model.
To make it easier understanding, here’s some data:
With ToTensor():
    [[0.8824, 0.8824, 0.9059,  ..., 0.8980, 0.8784, 0.8784],
     [0.8824, 0.9059, 0.8941,  ..., 0.8941, 0.8863, 0.8784],
     [0.8784, 0.8824, 0.8706,  ..., 0.8902, 0.8902, 0.8902],
     ...,
     [0.8941, 0.8941, 0.8980,  ..., 0.8941, 0.9098, 0.9216],
     [0.9059, 0.9020, 0.9059,  ..., 0.9098, 0.9255, 0.9255],
     [0.9059, 0.9059, 0.9059,  ..., 0.9216, 0.9255, 0.9255]],

    [[0.8902, 0.8902, 0.8941,  ..., 0.8588, 0.8627, 0.8627],
     [0.8902, 0.8784, 0.8824,  ..., 0.8549, 0.8588, 0.8627],
     [0.8706, 0.8706, 0.8784,  ..., 0.8510, 0.8588, 0.8667],
     ...,
     [0.9098, 0.8980, 0.8902,  ..., 0.8824, 0.8784, 0.8784],
     [0.8863, 0.8863, 0.8824,  ..., 0.8784, 0.8745, 0.8706],
     [0.8863, 0.8863, 0.8784,  ..., 0.8784, 0.8706, 0.8706]]])

With transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])
    [[2.0084, 2.0084, 2.0084,  ..., 1.9384, 1.9559, 1.9559],
     [2.0084, 1.9734, 1.9559,  ..., 1.9209, 1.9559, 1.9559],
     [1.9734, 1.9734, 1.8859,  ..., 1.9559, 2.0084, 2.0084],
     ...,
     [2.0259, 1.9734, 1.9909,  ..., 1.9909, 2.0434, 2.1134],
     [2.0084, 1.9909, 2.0259,  ..., 2.0434, 2.0959, 2.0959],
     [2.0084, 2.0084, 2.0259,  ..., 2.0434, 2.0959, 2.0959]],

    [[2.1868, 2.1868, 2.3088,  ..., 2.0125, 2.0474, 2.0474],
     [2.1868, 2.1694, 2.2217,  ..., 2.0474, 2.0648, 2.0474],
     [2.1520, 2.1520, 2.2043,  ..., 2.0823, 2.0997, 2.1346],
     ...,
     [2.1346, 2.1520, 2.1868,  ..., 2.1171, 2.1171, 2.1171],
     [2.1694, 2.1694, 2.2043,  ..., 2.1520, 2.1346, 2.1171],
     [2.1694, 2.1694, 2.2391,  ..., 2.1520, 2.1171, 2.1171]]])