diff --git a/chroma_NAG.py b/chroma_NAG.py
index 4e7b52e..5680464 100644
--- a/chroma_NAG.py
+++ b/chroma_NAG.py
@@ -213,7 +213,7 @@ class ChromaNAG:
         model_clone = model.clone()
         diffusion_model = model_clone.get_model_object("diffusion_model")
         diffusion_model.txt_in.to(device)
-        txt = diffusion_model.txt_in(nag_context.to(device, diffusion_model.txt_in.dtype))
+        txt = diffusion_model.txt_in(nag_context.to(device, mm.text_encoder_dtype()))
 
         # Chroma models have `double_blocks` where image and text tokens interact.
         # This is the equivalent of a cross-attention stage in other models.