Write the current version into model configs (#160)

neuralmagic · Sep 24, 2024 · 52afd24 · 52afd24
1 parent 799be8d
commit 52afd24
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 0 deletions.
diff --git a/src/compressed_tensors/__init__.py b/src/compressed_tensors/__init__.py
@@ -19,3 +19,4 @@
 from .config import *
 from .quantization import QuantizationConfig, QuantizationStatus
 from .utils import *
+from .version import *
diff --git a/src/compressed_tensors/compressors/model_compressor.py b/src/compressed_tensors/compressors/model_compressor.py
@@ -22,6 +22,7 @@
 
 import torch
 import transformers
+import compressed_tensors
 from compressed_tensors.base import (
     COMPRESSION_CONFIG_NAME,
     QUANTIZATION_CONFIG_NAME,
@@ -368,6 +369,7 @@ def update_config(self, save_directory: str):
             config_data[COMPRESSION_CONFIG_NAME][
                 SPARSITY_CONFIG_NAME
             ] = sparsity_config_data
+        config_data[COMPRESSION_CONFIG_NAME]["version"] = compressed_tensors.__version__
 
         with open(config_file_path, "w") as config_file:
             json.dump(config_data, config_file, indent=2, sort_keys=True)