make serializable

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

Author: kylesayrs

src/compressed_tensors/transform/factory/base.py

@@ -112,7 +112,7 @@ def _apply_to_module(self, module: Module, args: TransformArgs):
                 )
 
         # 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)
         self.transforms.append(transform)
         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