Fix :top_k limit in classification servings

elixir-nx · Sep 18, 2023 · 355ad10 · 355ad10
1 parent 7430d64
commit 355ad10
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/lib/bumblebee/text/text_classification.ex b/lib/bumblebee/text/text_classification.ex
@@ -36,7 +36,7 @@ defmodule Bumblebee.Text.TextClassification do
     scores_fun = fn params, input ->
       outputs = predict_fun.(params, input)
       scores = Shared.logits_to_scores(outputs.logits, scores_function)
-      k = min(top_k, Nx.size(scores))
+      k = min(top_k, Nx.axis_size(scores, 1))
       {top_scores, top_indices} = Nx.top_k(scores, k: k)
       {top_scores, top_indices}
     end

diff --git a/lib/bumblebee/vision/image_classification.ex b/lib/bumblebee/vision/image_classification.ex
@@ -40,7 +40,7 @@ defmodule Bumblebee.Vision.ImageClassification do
       input = Bumblebee.Featurizer.process_batch(featurizer, input)
       outputs = predict_fun.(params, input)
       scores = Shared.logits_to_scores(outputs.logits, scores_function)
-      k = min(top_k, Nx.size(scores))
+      k = min(top_k, Nx.axis_size(scores, 1))
       {top_scores, top_indices} = Nx.top_k(scores, k: k)
       {top_scores, top_indices}
     end