Unverified Commit 534512be authored by Mishig Davaadorj's avatar Mishig Davaadorj Committed by GitHub
Browse files

[flax] 'dtype' should not be part of self._internal_dict (#609)

parent 4b8880a3
...@@ -456,6 +456,9 @@ def flax_register_to_config(cls): ...@@ -456,6 +456,9 @@ def flax_register_to_config(cls):
# Make sure init_kwargs override default kwargs # Make sure init_kwargs override default kwargs
new_kwargs = {**default_kwargs, **init_kwargs} new_kwargs = {**default_kwargs, **init_kwargs}
# dtype should be part of `init_kwargs`, but not `new_kwargs`
if "dtype" in new_kwargs:
new_kwargs.pop("dtype")
# Get positional arguments aligned with kwargs # Get positional arguments aligned with kwargs
for i, arg in enumerate(args): for i, arg in enumerate(args):
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment