From 892078a0c946c90693c4c912cc1b7fdba297fc2a Mon Sep 17 00:00:00 2001 From: Maanu Grover Date: Wed, 11 Feb 2026 15:05:32 -0800 Subject: [PATCH] avoid direct pickle import Signed-off-by: Maanu Grover --- model_provider.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/model_provider.py b/model_provider.py index f8f6ccae01c..0c80c54dfdb 100644 --- a/model_provider.py +++ b/model_provider.py @@ -50,13 +50,9 @@ def model_provider( def oom_observer(device, alloc, device_alloc, device_free): # snapshot right after an OOM happened print('saving allocated state during OOM') - snapshot = torch.cuda.memory._snapshot() - from pickle import dump - dump( - snapshot, - open(f"oom_rank-{torch.distributed.get_rank()}_{args.memory_snapshot_path}", 'wb'), - ) + filename = f"oom_rank-{torch.distributed.get_rank()}_{args.memory_snapshot_path}" + torch.cuda.memory._dump_snapshot(filename) torch._C._cuda_attach_out_of_memory_observer(oom_observer)