Simplify saving model

awq/models/base.py

@@ -55,7 +55,9 @@ class BaseAWQForCausalLM(nn.Module):
         pass
 
     def save_quantized(self, save_dir, safetensors=False, shard_size="10GB"):
-        def _save_files(save_dir, model_name='', search_result=None):
+        save_dir = save_dir[:-1] if save_dir[-1] == '/' else save_dir
+
+        # Save model
         class EmptyModule(nn.Module):
             def __init__(self): super(EmptyModule, self).__init__()
             def forward(self, x): return x
@@ -66,9 +68,6 @@ class BaseAWQForCausalLM(nn.Module):
         # Remove empty state dict
         os.remove(f'{save_dir}/pytorch_model.bin')
 
-            if search_result is not None:
-                torch.save(search_result, f'{save_dir}/{model_name}')
-            else:
         # model_name has no extension, add it when saving state_dict
         model_name = 'model.safetensors' if safetensors else 'pytorch_model.bin'
 
@@ -96,14 +95,6 @@ class BaseAWQForCausalLM(nn.Module):
         with open(f'{save_dir}/quant_config.json', 'w+') as file:
             file.write(json.dumps(self.quant_config, indent=4))
         
-        save_dir = save_dir[:-1] if save_dir[-1] == '/' else save_dir
-
-        # Save model
-        if self.search_result is None or self.is_quantized:
-            _save_files(save_dir, '', search_result=None)
-        else:
-            model_name = 'awq_model_search_result.pt'
-            _save_files(save_dir, model_name, self.search_result)
         
     @classmethod
     def from_pretrained(self, model_path, model_type, torch_dtype: torch.dtype = torch.float16,