remove of unsqueece

maxkeller321 · maxkeller321 · commit c5160062363c · 2021-04-24T10:02:15.000+02:00
diff --git a/modAL/dropout.py b/modAL/dropout.py
@@ -318,7 +318,7 @@ def get_predictions(classifier: BaseEstimator, X: modALinput, dropout_layer_inde
                 logits = classifier.estimator.infer(samples)
                 prediction = logits_adaptor(logits, samples)
                 mask = ~prediction.isnan()
-                prediction[mask] = prediction[mask].unsqueeze(0).softmax(1)
+                prediction[mask] = prediction[mask].softmax(1)
                 probas.append(prediction)
         
         probas = torch.cat(probas)
