Convert TF tensor to Numpy array

bhky · Dec 1, 2023 · 4f98dda · 4f98dda
1 parent 86c0a81
commit 4f98dda
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/opennsfw2/_inference.py b/opennsfw2/_inference.py
@@ -61,7 +61,7 @@ def _predict_images_in_batches(
     prediction_batches: List[NDFloat32Array] = []
     for i in range(0, len(images), batch_size):
         batch = np.array(images[i: i + batch_size])
-        prediction_batches.append(model_(batch))
+        prediction_batches.append(model_(batch).numpy())
     predictions = np.concatenate(prediction_batches, axis=0)
     return predictions