From 063cdcfbf7dfb507250f76ff03e026c2bea455b6 Mon Sep 17 00:00:00 2001 From: Xuehai Pan Date: Thu, 6 Feb 2025 02:43:15 +0800 Subject: [PATCH] Update (base update) [ghstack-poisoned] --- torch/serialization.py | 6 +----- torchgen/local.py | 7 ++++--- 2 files changed, 5 insertions(+), 8 deletions(-) diff --git a/torch/serialization.py b/torch/serialization.py index 514cd318c1c..bff322e92f8 100644 --- a/torch/serialization.py +++ b/torch/serialization.py @@ -1340,11 +1340,7 @@ def load( ... weights_only=True, ... ) # type: ignore[attr-defined] # Map tensors from GPU 1 to GPU 0 - >>> torch.load( - ... "tensors.pt", - ... map_location={"cuda:1": "cuda:0"}, - ... weights_only=True, - ... ) + >>> torch.load("tensors.pt", map_location={"cuda:1": "cuda:0"}, weights_only=True) # Load tensor from io.BytesIO object # Loading from a buffer setting weights_only=False, warning this can be unsafe >>> with open("tensor.pt", "rb") as f: diff --git a/torchgen/local.py b/torchgen/local.py index 41f78539ca2..19045f4a948 100644 --- a/torchgen/local.py +++ b/torchgen/local.py @@ -39,9 +39,10 @@ def use_const_ref_for_mutable_tensors() -> bool: def use_ilistref_for_tensor_lists() -> bool: - assert ( - _locals.use_ilistref_for_tensor_lists is not None - ), "need to initialize local.use_ilistref_for_tensor_lists with local.parametrize" + assert _locals.use_ilistref_for_tensor_lists is not None, ( + "need to initialize local.use_ilistref_for_tensor_lists with " + "local.parametrize" + ) return _locals.use_ilistref_for_tensor_lists