Fix warning

Author: jonatanklosko

lib/bumblebee/conversion/pytorch_params.ex

@@ -84,12 +84,16 @@ defmodule Bumblebee.Conversion.PyTorchParams do
 
   defp state_dict?(%{} = dict) when not is_struct(dict) do
     Enum.all?(dict, fn {key, value} ->
-      is_binary(key) and Nx.LazyContainer.impl_for(value) != nil
+      is_binary(key) and implements_lazy_container?(value)
     end)
   end
 
   defp state_dict?(_other), do: false
 
+  defp implements_lazy_container?(value) do
+    Nx.LazyContainer.impl_for(value) != Nx.LazyContainer.Any or Nx.Container.impl_for(value) != nil
+  end
+
   defp init_params(model, params_expr, pytorch_state, params_mapping) do
     layers =
       model