From 534512bedb70ab6074fd4399a3ce15afc4394a37 Mon Sep 17 00:00:00 2001
From: Mishig Davaadorj <mishig.davaadorj@coloradocollege.edu>
Date: Thu, 22 Sep 2022 11:46:31 +0200
Subject: [PATCH] [flax] 'dtype' should not be part of self._internal_dict
 (#609)

---
 src/diffusers/configuration_utils.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/src/diffusers/configuration_utils.py b/src/diffusers/configuration_utils.py
index 7c9e4e46..19f58fd8 100644
--- a/src/diffusers/configuration_utils.py
+++ b/src/diffusers/configuration_utils.py
@@ -456,6 +456,9 @@ def flax_register_to_config(cls):
 
         # Make sure init_kwargs override default 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
         for i, arg in enumerate(args):