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)