Hello,
I am trying to implement an AL pipeline on my project using BALD query strategy. I am using this 15 lovely project but I am facing some issues, mainly because my project is a 2 class segmentation problem.
My input images are  RGB 256x256.
Following the project:
Initially, the author creates a dataloader, runs the prediction and applies softmax on top of those predictions.
def predict_prob_dropout_split(self, X, Y, n_drop):
        loader_te = DataLoader(self.handler(X, Y, transform=self.args['transform']),
                            shuffle=False, **self.args['loader_te_args'])

        self.clf.train()
        probs = torch.zeros([n_drop, len(Y), len(np.unique(Y))])
        for i in range(n_drop):
            print('n_drop {}/{}'.format(i+1, n_drop))
            with torch.no_grad():
                for x, y, idxs in loader_te:
                    x, y = x.to(self.device), y.to(self.device)
                    out, e1 = self.clf(x)
                    probs[i][idxs] += F.softmax(out, dim=1).cpu()
        
        return probs

However, I am just using a test dataset that I created using SubsetRandomSampler to generate a random subset of images to create my pool_loader.
pool_loader = DataLoader(dataset, batch_size=1, num_workers=num_workers,
                                          sampler=SubsetRandomSampler(pool_idx))

Then, I run my model in this pool_loader and the idea is to save the predictions as the author did, with the probs tensor in the line: probs[i][idxs] += F.softmax(out, dim=1).cpu()
pool_size = 10
probs = torch.zeros([10, pool_size, 2])
for i in range(10):
    with torch.no_grad():
        with tqdm(enumerate(pool_loader)) as iterator:
            for idx,batch in iterator:
                logits = model(batch.to(device, dtype=torch.float)).cpu().detach()
                logits = logits.squeeze()
                probs[i][idx] += logits
   

I replaced the len(np.unique(Y)) with 2, since I only have two classes.
logits.shape = torch.Size([1, 1, 256, 256])
After squeezing:
logits.shape = torch.Size([256, 256])
I think everything looks good until here but then I get the following error on line probs[i][idx] = logits

RuntimeError: expand(torch.FloatTensor{[256, 256]}, size=[2]): the number of sizes provided (1) must be greater or equal to the number of dimensions in the tensor (2)

The idea is to save the probabilities over those 10 runs for the 10 images of my pool loader. With this, calculate the entropy by doing this:
pb = probs.mean(0)
entropy1 = (-pb*torch.log(pb)).sum(1)
entropy2 = (-probs*torch.log(probs)).sum(2).mean(0)
U = entropy2 - entropy1
return idxs_unlabeled[U.sort()[1][:n]]

For me, it looks a very interesting use case and I would like to learn a bit more about it. First, I would be glad to know what I am doing wrong and understand how I can fit my results to calculate the entropy of my 10 trials.
Any suggestion is welcome 
Kind regards