Hello, I am running a Unet model with sigmoid as activation function and I am trying to get the softmax probabilites for each class. However, I am facing two problems:
First, the result of the softmax probability is always 1

tensor([[[[1., 1., 1.,  ..., 1., 1., 1.],
          [1., 1., 1.,  ..., 1., 1., 1.],
          [1., 1., 1.,  ..., 1., 1., 1.],
          ...,
          [1., 1., 1.,  ..., 1., 1., 1.],
          [1., 1., 1.,  ..., 1., 1., 1.],
          [1., 1., 1.,  ..., 1., 1., 1.]]],


        [[[1., 1., 1.,  ..., 1., 1., 1.],
          [1., 1., 1.,  ..., 1., 1., 1.],
          [1., 1., 1.,  ..., 1., 1., 1.],
          ...,
          [1., 1., 1.,  ..., 1., 1., 1.],
          [1., 1., 1.,  ..., 1., 1., 1.],
          [1., 1., 1.,  ..., 1., 1., 1.]]]])
tensor([[[[1., 1., 1.,  ..., 1., 1., 1.],
          [1., 1., 1.,  ..., 1., 1., 1.],
          [1., 1., 1.,  ..., 1., 1., 1.],
          ...,
          [1., 1., 1.,  ..., 1., 1., 1.],
          [1., 1., 1.,  ..., 1., 1., 1.],
          [1., 1., 1.,  ..., 1., 1., 1.]]],

Then, when trying to select the top two classes confidences for each sample,
toptwo = torch.topk(probabilities, 2, dim=1)[0]
I get
RuntimeError: selected index k out of range
Full code:
with torch.no_grad():
        
        with tqdm(enumerate(pool_loader)) as iterator:
            for i,batch in iterator:
                logits = model.forward(batch.to(device, dtype=torch.float)).cpu().detach()
                probabilities = F.softmax(logits, dim=1)
                print(probabilities)
            
            # Select the top two class confidences for each sample
            toptwo = torch.topk(probabilities, 2, dim=1)[0]

I suspect I am messing up somewhere in the batch, so I will leave here how my 2d batch looks like:
batch[1]
tensor([[[ 1.7694,  1.8379,  1.7865,  ...,  1.8208,  1.7865,  1.7865],
         [ 1.7523,  1.8550,  1.8208,  ...,  1.8208,  1.7694,  1.7865],
         [ 1.7694,  1.8379,  1.8208,  ...,  1.7865,  1.7865,  1.7865],
         ...,
         [-0.2171, -0.2342, -0.2171,  ...,  1.1700,  1.1015,  1.1015],
         [-0.1486, -0.2171, -0.1828,  ...,  1.0159,  1.0502,  0.9817],
         [-0.1486, -0.1657, -0.1486,  ...,  0.9988,  0.9817,  1.0159]],

        [[ 1.8683,  1.9209,  1.8683,  ...,  1.8683,  1.8859,  1.8859],
         [ 1.9034,  1.8859,  1.9034,  ...,  1.8859,  1.8859,  1.8859],
         [ 1.8683,  1.9034,  1.8859,  ...,  1.8859,  1.8859,  1.8859],
         ...,
         [-0.8452, -0.8627, -0.8277,  ...,  0.9230,  0.8354,  0.7654],
         [-0.7577, -0.8452, -0.8102,  ...,  0.8179,  0.7479,  0.6254],
         [-0.7052, -0.7402, -0.7577,  ...,  0.7654,  0.6954,  0.7129]],

        [[ 2.2391,  2.2740,  2.2391,  ...,  2.2391,  2.2391,  2.2566],
         [ 2.2391,  2.2740,  2.2566,  ...,  2.2391,  2.2391,  2.2391],
         [ 2.2217,  2.2566,  2.2566,  ...,  2.2391,  2.2566,  2.2217],
         ...,
         [ 0.6705,  0.6008,  0.6879,  ...,  1.8557,  1.8208,  1.7685],
         [ 0.8274,  0.7054,  0.7402,  ...,  1.8383,  1.8208,  1.7163],
         [ 0.8099,  0.7925,  0.7925,  ...,  1.8383,  1.8034,  1.8208]]],
       dtype=torch.float64)

batch[0]:
tensor([[[ 1.7694,  1.8379,  1.7865,  ...,  1.8208,  1.7865,  1.7865],
         [ 1.7523,  1.8550,  1.8208,  ...,  1.8208,  1.7694,  1.7865],
         [ 1.7694,  1.8379,  1.8208,  ...,  1.7865,  1.7865,  1.7865],
         ...,
         [-0.2171, -0.2342, -0.2171,  ...,  1.1700,  1.1015,  1.1015],
         [-0.1486, -0.2171, -0.1828,  ...,  1.0159,  1.0502,  0.9817],
         [-0.1486, -0.1657, -0.1486,  ...,  0.9988,  0.9817,  1.0159]],

        [[ 1.8683,  1.9209,  1.8683,  ...,  1.8683,  1.8859,  1.8859],
         [ 1.9034,  1.8859,  1.9034,  ...,  1.8859,  1.8859,  1.8859],
         [ 1.8683,  1.9034,  1.8859,  ...,  1.8859,  1.8859,  1.8859],
         ...,
         [-0.8452, -0.8627, -0.8277,  ...,  0.9230,  0.8354,  0.7654],
         [-0.7577, -0.8452, -0.8102,  ...,  0.8179,  0.7479,  0.6254],
         [-0.7052, -0.7402, -0.7577,  ...,  0.7654,  0.6954,  0.7129]],

        [[ 2.2391,  2.2740,  2.2391,  ...,  2.2391,  2.2391,  2.2566],
         [ 2.2391,  2.2740,  2.2566,  ...,  2.2391,  2.2391,  2.2391],
         [ 2.2217,  2.2566,  2.2566,  ...,  2.2391,  2.2566,  2.2217],
         ...,
         [ 0.6705,  0.6008,  0.6879,  ...,  1.8557,  1.8208,  1.7685],
         [ 0.8274,  0.7054,  0.7402,  ...,  1.8383,  1.8208,  1.7163],
         [ 0.8099,  0.7925,  0.7925,  ...,  1.8383,  1.8034,  1.8208]]],
       dtype=torch.float64)

Any help would be amazing, Ive been trying to understand what is wrong for hours.
Regards