Add BytesIO to safe_globals

examples/inference/gpt/gpt_dynamic_inference.py

@@ -60,11 +60,6 @@
 import io
 import megatron
 
-torch.serialization.add_safe_globals([io.BytesIO])
-torch.serialization.add_safe_globals([megatron.core.rerun_state_machine.RerunState])
-torch.serialization.add_safe_globals([megatron.core.rerun_state_machine.RerunDiagnostic])
-
-
 
 def add_dynamic_inference_args(parser: ArgumentParser) -> ArgumentParser:
     """Dynamic inference arguments."""

megatron/core/safe_globals.py

@@ -1,6 +1,7 @@
 # Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
 
 from argparse import Namespace
+from io import BytesIO
 from pathlib import PosixPath
 from types import SimpleNamespace
 
@@ -26,6 +27,7 @@
     RerunDiagnostic,
     RerunMode,
     RerunState,
+    BytesIO,
 ]