Update quantization/apply_gptq_save_marlin.py

quantization/apply_gptq_save_marlin.py

@@ -46,7 +46,7 @@ if __name__ == "__main__":
         device_map="auto")
     model.quantize(examples)
     
-    gptq_save_dir = args.gptq_save_dir
+    gptq_save_dir = "./tmp-gptq"
     print(f"Saving gptq model to {gptq_save_dir}")
     model.save_pretrained(gptq_save_dir)
     tokenizer.save_pretrained(gptq_save_dir)
@@ -55,14 +55,14 @@ if __name__ == "__main__":
     gc.collect()
 
     print("Reloading in marlin format")
-    gptq_save_dir = "./tmp-gptq"
+    
     marlin_model = AutoGPTQForCausalLM.from_quantized(
         gptq_save_dir, 
         use_marlin=True, 
         device_map="auto")
 
     print("Saving in marlin format")
-    marlin_model.save_pretrained(args.marlin_save_dir)
-    tokenizer.save_pretrained(args.marlin_save_dir)
+    marlin_model.save_pretrained(args.save_dir)
+    tokenizer.save_pretrained(args.save_dir)
 
     shutil.rmtree(gptq_save_dir)