diff --git a/HunYuanDiT/tenc.py b/HunYuanDiT/tenc.py
index 7883982..311543b 100644
--- a/HunYuanDiT/tenc.py
+++ b/HunYuanDiT/tenc.py
@@ -4,6 +4,7 @@ import torch
 from transformers import AutoTokenizer, modeling_utils
 from transformers import T5Config, T5EncoderModel, BertConfig, BertModel
 
+from comfy import model_management
 import comfy.model_patcher
 import comfy.utils