try_to_allocate_tensor_on_cpu

Author: maxkeller321

modAL/dropout.py

@@ -317,12 +317,11 @@ def get_predictions(classifier: BaseEstimator, X: modALinput, dropout_layer_inde
             logits = classifier.estimator.infer(samples)
             prediction = logits_adaptor(logits, samples)
 
-            if probas is None: probas = torch.empty((number_of_samples, prediction.shape[-1]), device=prediction.device)
-            probas[range(sample_per_forward_pass*index, sample_per_forward_pass*(index+1)), :] = prediction
+            mask = ~prediction.isnan()
+            prediction[mask] = prediction[mask].unsqueeze(0).softmax(1)
+            if probas is None: probas = torch.empty((number_of_samples, prediction.shape[-1]), device='cpu')
+            probas[range(sample_per_forward_pass*index, sample_per_forward_pass*(index+1)), :] = prediction.cpu()
 
-
-        mask = ~probas.isnan()
-        probas[mask] = probas[mask].unsqueeze(0).softmax(1)
         probas = to_numpy(probas)
         predictions.append(probas)