make serializable (#396)

Signed-off-by: Kyle Sayers <[email protected]>

Author: kylesayrs

src/compressed_tensors/transform/factory/base.py

@@ -99,7 +99,7 @@ def _apply_to_module(self, module: Module, args: TransformArgs):
         :param args: defines how the transform will be applied to the target module
         """
         # create transform as submodule
-        transform_name = f"{self.name}_{args.location.value}"
+        transform_name = f"{self.name}_{args.location}"
         transform = self.create_transform(module, args)
         register_offload_module(module, transform_name, transform)
 

src/compressed_tensors/transform/transform_args.py

@@ -46,7 +46,7 @@ class TransformLocation(str, Enum):
     Q_ATTN = "q_attn"
 
 
-class TransformArgs(BaseModel):
+class TransformArgs(BaseModel, use_enum_values=True):
     """
     Arguments which define *how* and where a transform should be applied to a model