From 4dbe5ef7e3fd80779dd94018a341f94cbd1b4098 Mon Sep 17 00:00:00 2001 From: Huanyu Zhang Date: Tue, 10 Dec 2024 09:09:05 -0800 Subject: [PATCH] Fix ``torch.load()`` in model_utils.py (#696) Summary: Pull Request resolved: https://github.com/pytorch/opacus/pull/696 Addressing Github issue #690 (https://github.com/pytorch/opacus/issues/690). Since Opacus is primarily intended for academic use cases, it is considered low risk to utilize torch.load() with the setting "weights_only=False". Additionally, reminders have been added to the function description. Reviewed By: iden-kalemaj Differential Revision: D66999322 fbshipit-source-id: 6eddc7b5a0390809ef25fc9fe9e5d28bf6b55130 --- opacus/utils/module_utils.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/opacus/utils/module_utils.py b/opacus/utils/module_utils.py index 28146cef0..ee0437a9d 100644 --- a/opacus/utils/module_utils.py +++ b/opacus/utils/module_utils.py @@ -89,6 +89,8 @@ def clone_module(module: nn.Module) -> nn.Module: """ Handy utility to clone an nn.Module. PyTorch doesn't always support copy.deepcopy(), so it is just easier to serialize the model to a BytesIO and read it from there. + When ``weights_only=False``, ``torch.load()`` uses "pickle" module implicity, which is known to be insecure. + Only load the model you trust. Args: module: The module to clone @@ -99,7 +101,7 @@ def clone_module(module: nn.Module) -> nn.Module: with io.BytesIO() as bytesio: torch.save(module, bytesio) bytesio.seek(0) - module_copy = torch.load(bytesio) + module_copy = torch.load(bytesio, weights_only=False) next_param = next( module.parameters(), None ) # Eg, InstanceNorm with affine=False has no params