commit 48505356f519d686bf86e1f32377dac22f4ebf4f
author Andrij David <andrijdavid@gmail.com> Tue Jun 14 16:04:45 2022 +0000
committer PyTorch MergeBot <pytorchmergebot@users.noreply.github.com> Tue Jun 14 16:04:45 2022 +0000
tree bdf8fca9299cf657bc1a907b6f827a916214fcf0
parent 9d351b3ddde9d9e7bb0976707601c023d5ab8a4b

Propagate map_location arg to torch.jit.load in torch.load (#78733)

Fixes #78331

Pull Request resolved: https://github.com/pytorch/pytorch/pull/78733
Approved by: https://github.com/davidberard98

torch/serialization.py

diff --git a/torch/serialization.py b/torch/serialization.py
index e778817..8262b96 100644
--- a/torch/serialization.py
+++ b/torch/serialization.py
@@ -718,7 +718,7 @@
                                   " dispatching to 'torch.jit.load' (call 'torch.jit.load' directly to"
                                   " silence this warning)", UserWarning)
                     opened_file.seek(orig_position)
-                    return torch.jit.load(opened_file)
+                    return torch.jit.load(opened_file, map_location=map_location)
                 return _load(opened_zipfile, map_location, pickle_module, **pickle_load_args)
         return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)