diff --git a/fooocus_version.py b/fooocus_version.py
index e83c613b..024361b7 100644
--- a/fooocus_version.py
+++ b/fooocus_version.py
@@ -1 +1 @@
-version = '2.0.3'
+version = '2.0.4'
diff --git a/modules/patch.py b/modules/patch.py
index 08474956..09ec6fc2 100644
--- a/modules/patch.py
+++ b/modules/patch.py
@@ -70,7 +70,8 @@ def sdxl_encode_adm_patched(self, **kwargs):
 
 
 def text_encoder_device_patched():
-    return torch.device("cpu")
+    # Fooocus's style system uses text encoder much more times than comfy so this makes things much faster.
+    return comfy.model_management.get_torch_device()
 
 
 def patch_all():