place multi-processing init to main method (#443)

NVIDIA · Feb 10, 2023 · 7634698 · 7634698
1 parent fff55e7
commit 7634698
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/examples/pytorch/gpt/utils/huggingface_gpt_convert.py b/examples/pytorch/gpt/utils/huggingface_gpt_convert.py
@@ -158,9 +158,7 @@ def split_and_convert(args):
         "mlp.dense_4h_to_h.bias",
         "mlp.dense_4h_to_h.weight",
     ]
-
-    torch.multiprocessing.set_start_method("spawn")
-    torch.multiprocessing.set_sharing_strategy("file_system")
+
     pool = multiprocessing.Pool(args.processes)
     for name, param in model.named_parameters():
         if name.find("weight") == -1 and name.find("bias") == -1:
@@ -201,4 +199,6 @@ def split_and_convert(args):
         print("{}: {}".format(key, vars(args)[key]))
     print("========================================")
 
+    torch.multiprocessing.set_start_method("spawn")
+    torch.multiprocessing.set_sharing_strategy("file_system")
     split_and_convert(args)