Use provided HF token explicitly for download and push

wave_tank_surgeon.py

@@ -607,16 +607,27 @@ def main():
     print(f"Training: {args.steps} steps, lr={args.lr}")
     print("=" * 70)
 
+    hf_token = (
+        os.environ.get('HF_TOKEN')
+        or os.environ.get('HUGGING_FACE_HUB_TOKEN')
+        or os.environ.get('HUGGINGFACEHUB_API_TOKEN')
+    )
+
     # Load model and tokenizer
     print(f"Loading model: {args.model}...")
     from transformers import AutoModelForCausalLM, AutoTokenizer
 
-    tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True)
+    tokenizer = AutoTokenizer.from_pretrained(
+        args.model,
+        trust_remote_code=True,
+        token=hf_token,
+    )
     model = AutoModelForCausalLM.from_pretrained(
         args.model,
         torch_dtype=torch_dtype,
         device_map='auto' if torch.cuda.is_available() else None,
         trust_remote_code=True,
+        token=hf_token,
     )
 
     if device.type == 'cpu' or (device.type == 'cuda' and not hasattr(model, 'hf_device_map')):
@@ -794,11 +805,11 @@ def main():
     # Push to Hub if requested
     if args.push_to_hub:
         print(f"Pushing to HuggingFace Hub: {args.push_to_hub}...")
-        model.push_to_hub(args.push_to_hub, private=args.private)
-        tokenizer.push_to_hub(args.push_to_hub, private=args.private)
+        model.push_to_hub(args.push_to_hub, private=args.private, token=hf_token)
+        tokenizer.push_to_hub(args.push_to_hub, private=args.private, token=hf_token)
         # Also push surgery report
         from huggingface_hub import HfApi
-        api = HfApi()
+        api = HfApi(token=hf_token)
         api.upload_file(
             path_or_fileobj=os.path.join(output_dir, 'surgery_report.json'),
             path_in_repo='surgery_report.json',