Load initial expert location metadata on CPU (#32435)

2026-07-26 20:44:25 +08:00
parent a76b74cbe0
commit e8a635a412
1 changed files with 1 additions and 1 deletions
@@ -773,7 +773,7 @@ def compute_initial_expert_location_metadata(

    # TODO unify with the utils function
    if data.endswith(".pt"):
-        data_dict = torch.load(data, weights_only=True)
+        data_dict = torch.load(data, weights_only=True, map_location="cpu")
    elif data.endswith(".json"):
        data_dict = json.loads(Path(data).read_text())
    else: