A vulnerability was identified in PyTorch 2.10.0. The affected element is an unknown function of the component pt2 Loading Handler. The manipulation leads to deserialization. The attack can only be pe

@@ -422,6 +422,7 @@ def serialize_torch_artifact(
  
 def deserialize_torch_artifact(
     serialized: dict[str, Any] | tuple[Any, ...] | bytes,
+    weights_only: bool = False,
 ):
     if isinstance(serialized, (dict, tuple)):
         return serialized
@@ -429,18 +430,22 @@ def deserialize_torch_artifact(
         return {}
     buffer = io.BytesIO(serialized)
     buffer.seek(0)
-    # weights_only=False as we want to load custom objects here (e.g. ScriptObject)
-    try:
+    
+    if weights_only:
         artifact = torch.load(buffer, weights_only=True)
-    except Exception as e:
-        buffer.seek(0)
-        artifact = torch.load(buffer, weights_only=False)
-        log.warning(
-            "Fallback to weights_only=False succeeded. "
-            "Loaded object of type %s after initial failure: %s",
-            type(artifact),
-            exc_info=e,
-        )
+    else:
+        # weights_only=False as we want to load custom objects here (e.g. ScriptObject)
+        try:
+            artifact = torch.load(buffer, weights_only=True)
+        except Exception as e:
+            buffer.seek(0)
+            artifact = torch.load(buffer, weights_only=False)
+            log.warning(
+                "Fallback to weights_only=False succeeded. "
+                "Loaded object of type %s after initial failure: %s",
+                type(artifact),
+                exc_info=e,
+            )
     if not isinstance(artifact, (tuple, dict)):
         raise AssertionError(f"expected tuple or dict, got {type(artifact).__name__}")
     return artifact
@@ -2830,6 +2835,8 @@ def deserialize(
         | bytes
         | None = None,
         symbol_name_to_range: dict[str, symbolic_shapes.ValueRanges] | None = None,
+        *,
+        weights_only: bool = False,
     ) -> Result:
         global _CURRENT_DESERIALIZER
         if _CURRENT_DESERIALIZER is not None:
@@ -2872,7 +2879,7 @@ def deserialize(
                 "Identity": torch.utils._sympy.functions.Identity,
             }
             self.symbol_name_to_symbol: dict[str, sympy.Symbol] = {}
-            self.constants = deserialize_torch_artifact(constants)
+            self.constants = deserialize_torch_artifact(constants, weights_only=weights_only)
             self.signature = self.deserialize_signature(
                 serialized_graph_module.signature
             )
@@ -2908,7 +2915,7 @@ def deserialize(
                 self.shape_env.unbacked_symint_counter += 1
  
             if example_inputs is not None and len(example_inputs) > 0:
-                self.example_inputs = deserialize_torch_artifact(example_inputs)
+                self.example_inputs = deserialize_torch_artifact(example_inputs, weights_only=weights_only)
             else:
                 self.example_inputs = None
             self.deserialize_graph(serialized_graph_module.graph)
@@ -2934,7 +2941,7 @@ def deserialize(
                 signature=self.signature,
                 module_call_graph=module_call_graph,
                 names_to_symbols=self.symbol_name_to_symbol,
-                state_dict=deserialize_torch_artifact(serialized_state_dict),
+                state_dict=deserialize_torch_artifact(serialized_state_dict, weights_only=weights_only),
                 constants=self.constants,
                 example_inputs=self.example_inputs,
             )
@@ -3494,6 +3501,7 @@ def deserialize(
         | bytes
         | None = None,
         *,
+        weights_only: bool = False,
         _unsafe_skip_version_check=False,
     ) -> ep.ExportedProgram:
         if not isinstance(exported_program, ExportedProgram):
@@ -3525,6 +3533,7 @@ def deserialize(
             constants,
             example_inputs,
             symbol_name_to_range,
+            weights_only=weights_only,
         )
         range_constraints = self.deserialize_range_constraints(
             symbol_name_to_range,
@@ -3687,6 +3696,7 @@ def deserialize(
     artifact: SerializedArtifact,
     expected_opset_version: dict[str, int] | None = None,
     *,
+    weights_only: bool = False,
     _unsafe_skip_version_check=False,
 ) -> ep.ExportedProgram:
     if not isinstance(artifact.exported_program, bytes):
@@ -3701,6 +3711,7 @@ def deserialize(
         artifact.state_dict,
         artifact.constants,
         artifact.example_inputs,
+        weights_only=weights_only,
         _unsafe_skip_version_check=_unsafe_skip_version_check,
     )
  

@@ -285,6 +285,7 @@ def load(
     *,
     extra_files: dict[str, Any] | None = None,
     expected_opset_version: dict[str, int] | None = None,
+    weights_only: bool = False,
 ) -> ExportedProgram:
     """
  
@@ -309,6 +310,11 @@ def load(
         expected_opset_version (Optional[Dict[str, int]]): A map of opset names
          to expected opset versions
  
+        weights_only (bool): Indicates whether to only load weights. If False,
+         it allows loading custom objects and other complex artifacts, which
+         could be a security risk if the file is from an untrusted source.
+         Defaults to False for backward compatibility.
+ 
     Returns:
         An :class:`ExportedProgram` object
  
@@ -343,6 +349,7 @@ def load(
         pt2_contents = load_pt2(
             f,
             expected_opset_version=expected_opset_version,
+            weights_only=weights_only,
         )
     except RuntimeError:
         log.warning("Ran into the following error when deserializing", exc_info=True)
@@ -429,7 +436,11 @@ def load(
         )
  
         # Deserialize ExportedProgram
-        ep = deserialize(artifact, expected_opset_version)
+        ep = deserialize(
+            artifact,
+            expected_opset_version,
+            weights_only=weights_only,
+        )
  
         return ep
  

@@ -845,6 +845,7 @@ def _load_payload_config(
 def _load_state_dict(
     archive_reader: PT2ArchiveReader,
     model_name: str,
+    weights_only: bool = False,
 ) -> dict[str, torch.Tensor] | bytes:
     # Make it BC compatible with legacy weight files
     legacy_weights_file = f"{WEIGHTS_DIR}{model_name}.pt"
@@ -872,7 +873,7 @@ def _load_state_dict(
                     os.path.join(WEIGHTS_DIR, payload_meta.path_name)
                 )
                 state_dict[weight_fqn] = torch.load(
-                    io.BytesIO(weight_bytes), weights_only=False
+                    io.BytesIO(weight_bytes), weights_only=weights_only
                 )
             else:
                 tensor_meta = payload_meta.tensor_meta
@@ -901,6 +902,7 @@ def _load_state_dict(
 def _load_constants(
     archive_reader: PT2ArchiveReader,
     model_name: str,
+    weights_only: bool = False,
 ) -> dict[str, torch.Tensor] | bytes:
     # Make it BC compatible with legacy constant files
     legacy_constants_file = f"{CONSTANTS_DIR}{model_name}.pt"
@@ -930,7 +932,7 @@ def _load_constants(
                         os.path.join(CONSTANTS_DIR, path_name)
                     )
                     constants[constant_fqn] = torch.load(
-                        io.BytesIO(constant_bytes), weights_only=False
+                        io.BytesIO(constant_bytes), weights_only=weights_only
                     )
                 else:
                     tensor_meta = payload_meta.tensor_meta
@@ -949,6 +951,10 @@ def _load_constants(
                     constants[constant_fqn] = constant_tensor
  
             elif path_name.startswith(CUSTOM_OBJ_FILENAME_PREFIX):
+                if weights_only:
+                    raise RuntimeError(
+                        f"Unsupported constant type {path_name} when weights_only is True"
+                    )
                 constant_bytes = archive_reader.read_bytes(
                     os.path.join(CONSTANTS_DIR, path_name)
                 )
@@ -964,6 +970,7 @@ def _load_exported_programs(
     archive_reader: PT2ArchiveReader,
     file_names: list[str],
     expected_opset_version: dict[str, int] | None,
+    weights_only: bool = False,
 ) -> dict[str, ExportedProgram]:
     exported_program_files = [
         file for file in file_names if file.startswith(MODELS_DIR)
@@ -986,8 +993,8 @@ def _load_exported_programs(
         serialized_exported_program = _bytes_to_dataclass(
             schema.ExportedProgram, exported_program_bytes
         )
-        state_dict = _load_state_dict(archive_reader, model_name)
-        constants = _load_constants(archive_reader, model_name)
+        state_dict = _load_state_dict(archive_reader, model_name, weights_only=weights_only)
+        constants = _load_constants(archive_reader, model_name, weights_only=weights_only)
  
         ep = ExportedProgramDeserializer(expected_opset_version).deserialize(
             serialized_exported_program,
@@ -1062,6 +1069,7 @@ def load_pt2(
     num_runners: int = 1,
     device_index: int = -1,
     load_weights_from_disk: bool = False,
+    weights_only: bool = False,
 ) -> PT2ArchiveContents:  # type: ignore[type-arg]
     """
     Loads all the artifacts previously saved with ``package_pt2``.
@@ -1118,7 +1126,7 @@ def load_pt2(
         file_names = archive_reader.get_file_names()
  
         exported_programs = _load_exported_programs(
-            archive_reader, file_names, expected_opset_version
+            archive_reader, file_names, expected_opset_version, weights_only=weights_only
         )
         extra_files = _load_extra_files(archive_reader, file_names)
  
@@ -1144,7 +1152,7 @@ def load_pt2(
                     len(WEIGHTS_DIR) :
                 ]  # remove data/weights/ prefix
                 weight_bytes = archive_reader.read_bytes(file)
-                loaded_weight = torch.load(io.BytesIO(weight_bytes))
+                loaded_weight = torch.load(io.BytesIO(weight_bytes), weights_only=weights_only)
                 weights[weight_file_name] = loaded_weight
  
     if isinstance(f, (io.IOBase, IO)):