From c97cfd1e83490e97a5c24de8a61e9b30b3cedd1f Mon Sep 17 00:00:00 2001
From: Kohaku-Blueleaf <59680068+KohakuBlueleaf@users.noreply.github.com>
Date: Sat, 9 Mar 2024 15:33:58 +0800
Subject: [PATCH] Fix dtype mismatch

---
 scripts/tilevae.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/scripts/tilevae.py b/scripts/tilevae.py
index fea1306..935e429 100644
--- a/scripts/tilevae.py
+++ b/scripts/tilevae.py
@@ -508,6 +508,7 @@ class VAEHook:
         @return: image
         """
         device = next(self.net.parameters()).device
+        dtype = next(self.net.parameters()).dtype
         net = self.net
         tile_size = self.tile_size
         is_decoder = self.is_decoder
@@ -647,7 +648,7 @@ class VAEHook:
 
         # Done!
         pbar.close()
-        return result if result is not None else result_approx.to(device)
+        return result.to(dtype) if result is not None else result_approx.to(device, dtype=dtype)
 
 
 class Script(scripts.Script):