From 8e3b3ba9841344ac7d27589903c90605d55c56db Mon Sep 17 00:00:00 2001 From: Antonio Bellotta Date: Wed, 21 Feb 2024 19:54:51 +0100 Subject: [PATCH] Use run_only_rank0 to broadcast allocation data --- neurodamus/utils/memory.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/neurodamus/utils/memory.py b/neurodamus/utils/memory.py index 3bc4e9895..b6d9ddc65 100644 --- a/neurodamus/utils/memory.py +++ b/neurodamus/utils/memory.py @@ -200,7 +200,8 @@ def export_allocation_stats(rank_allocation, filename): f.write(compressed_data) -def import_allocation_stats(filename): +@run_only_rank0 +def import_allocation_stats(filename) -> dict: """ Import allocation dictionary from serialized pickle file. """