Use rank instead of size

Co-authored-by: Mateusz Sluszniak <[email protected]>
elixir-nx · Oct 16, 2023 · 88582f2 · 88582f2
1 parent 92710d6
commit 88582f2
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/lib/scholar/options.ex b/lib/scholar/options.ex
@@ -102,7 +102,7 @@ defmodule Scholar.Options do
   end
 
   def beta(beta) do
-    if (is_number(beta) and beta >= 0) or (Nx.is_tensor(beta) and Nx.size(beta) == 1) do
+    if (is_number(beta) and beta >= 0) or (Nx.is_tensor(beta) and Nx.rank(beta) == 0) do
       {:ok, beta}
     else
       {:error, "expect 'beta' to be in the range [0, inf]"}