import ModelCompressor for type checking (#220)

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>
neuralmagic · Jan 23, 2025 · c026704 · c026704
1 parent fc711ca
commit c026704
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/src/compressed_tensors/utils/helpers.py b/src/compressed_tensors/utils/helpers.py
@@ -14,13 +14,17 @@
 
 import warnings
 from functools import wraps
-from typing import Any, Callable, Dict, List, Optional
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional
 
 import numpy
 import torch
 from transformers import AutoConfig
 
 
+if TYPE_CHECKING:
+    from compressed_tensors.compressors import ModelCompressor
+
+
 __all__ = [
     "infer_compressor_from_model_config",
     "fix_fsdp_module_name",