Update hf_merge.py

hf_merge.py

@@ -5,6 +5,7 @@ import requests
 from tqdm import tqdm
 from huggingface_hub import HfApi, hf_hub_download
 from merge import merge_folder, map_tensors_to_files, copy_nontensor_files, save_tensor_map
+import torch
 import logging
 
 # Set up logging
@@ -92,8 +93,10 @@ class ModelMerger:
         repo_manager.delete_repo(repo_path)
         repo_manager.download_repo(repo_name, repo_path)
 
+        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
         try:
-            self.tensor_map = merge_folder(self.tensor_map, repo_path, p, lambda_val)
+            self.tensor_map = merge_folder(self.tensor_map, repo_path, p, lambda_val, device)
             logging.info(f"Merged {repo_name}")
         except Exception as e:
             logging.error(f"Error merging {repo_name}: {e}")