Oh, sorry. I will try to explain then.
My input is 256x256 images and this is a segmentation problem. I implementing an active learning pipeline, therefore I am trying to get the probabilities and then get the top two classes confidences for each sample. Basically, the margin sampling query strategy.
I created a pool loader, that samples my test_dataset with random indices:
pool_idx = random.sample(range(1, len(test_dataset)), pool_size)
    pool_loader = DataLoader(dataset, batch_size=batch_size, num_workers=num_workers,
                                          #sampler=SubsetRandomSampler(unlabeled_idx[pool_idx]))
                                          sampler=SubsetRandomSampler(pool_idx))

Where the batch size is equal to 2.
Then I create an iterator to go through my pool_loader and get the prediction for each sample:


with tqdm(enumerate(test_loader)) as iterator:
    for i,batch in iterator:
        outputs=model.forward(batch.to(DEVICE, dtype=torch.float)).cpu().detach()

And the result is what I told you, the 0.0039 tensor.
I just want the softmax with the probabilities for each class and then calculate the margin.
Note that the activation function of my model (Unet) is sigmoid. I am afraid that applying softmax on top of sigmoid might be complicating things a bit. This is also a two class problem (Positive/Negative).
I hope this gives a better comprehension of my problem.
Thank you very much for your help!