From 077d704658218c3b6a0d6689ad471cf55075a2a3 Mon Sep 17 00:00:00 2001 From: jingxuz Date: Tue, 22 Jul 2025 16:04:46 +0800 Subject: [PATCH] fix(torch): add barrier before process group destruction in training example --- advanced/pytorch-example/main.py | 1 + 1 file changed, 1 insertion(+) diff --git a/advanced/pytorch-example/main.py b/advanced/pytorch-example/main.py index 5a64c11..4a2c68a 100644 --- a/advanced/pytorch-example/main.py +++ b/advanced/pytorch-example/main.py @@ -84,6 +84,7 @@ def transform(example): torch.save(model.module.state_dict(), "mnist_model.pth") print("Model saved as mnist_model.pth") + dist.barrier() dist.destroy_process_group() if __name__ == "__main__":