diff --git a/test/ex_vision/style_transfer/style_transfer_test.exs b/test/ex_vision/style_transfer/style_transfer_test.exs
index 45cbd71..21e7c9b 100644
--- a/test/ex_vision/style_transfer/style_transfer_test.exs
+++ b/test/ex_vision/style_transfer/style_transfer_test.exs
@@ -55,6 +55,7 @@ for {module, opts} <- TestConfiguration.configuration() do
       diff_sum =
         expected_result
         |> Nx.equal(result)
+        |> Nx.as_type(:u64)
         |> Nx.reduce(0, fn x, y -> Nx.add(x, y) end)
         # |> Nx.reduce_max()
         |> Nx.to_number()
@@ -63,6 +64,7 @@ for {module, opts} <- TestConfiguration.configuration() do
         expected_result
         |> Nx.subtract(result)
         |> Nx.abs()
+        |> Nx.as_type(:u64)
         |> Nx.reduce(0, fn x, y -> Nx.add(x, y) end)
         # |> Nx.reduce_max()
         |> Nx.to_number()