Autofix

Signed-off-by: Eric Kerfoot <[email protected]>

Author: ericspod

monai/data/__init__.py

@@ -77,8 +77,7 @@
 from .test_time_augmentation import TestTimeAugmentation
 from .thread_buffer import ThreadBuffer, ThreadDataLoader
 from .torchscript_utils import load_net_with_metadata, save_net_with_metadata
-from .utils import (
-    # PICKLE_KEY_SUFFIX,
+from .utils import (  # PICKLE_KEY_SUFFIX,
     affine_to_spacing,
     compute_importance_map,
     compute_shape_offset,

monai/data/dataset.py

@@ -12,9 +12,7 @@
 from __future__ import annotations
 
 import collections.abc
-from io import BytesIO
 import math
-from pickle import UnpicklingError
 import shutil
 import sys
 import tempfile
@@ -23,9 +21,11 @@
 import warnings
 from collections.abc import Callable, Sequence
 from copy import copy, deepcopy
+from io import BytesIO
 from multiprocessing.managers import ListProxy
 from multiprocessing.pool import ThreadPool
 from pathlib import Path
+from pickle import UnpicklingError
 from typing import IO, TYPE_CHECKING, Any, cast
 
 import numpy as np
@@ -254,7 +254,7 @@ def __init__(
                 this arg is used by `torch.save`, for more details, please check:
                 https://pytorch.org/docs/stable/generated/torch.save.html#torch.save,
                 and ``monai.data.utils.SUPPORTED_PICKLE_MOD``.
-            pickle_protocol: specifies pickle protocol when saving, with `torch.save`. 
+            pickle_protocol: specifies pickle protocol when saving, with `torch.save`.
                 Defaults to torch.serialization.DEFAULT_PROTOCOL. For more details, please check:
                 https://pytorch.org/docs/stable/generated/torch.save.html#torch.save.
             hash_transform: a callable to compute hash from the transform information when caching.
@@ -461,7 +461,7 @@ def __init__(
                 this arg is used by `torch.save`, for more details, please check:
                 https://pytorch.org/docs/stable/generated/torch.save.html#torch.save,
                 and ``monai.data.utils.SUPPORTED_PICKLE_MOD``.
-            pickle_protocol: specifies pickle protocol when saving, with `torch.save`. 
+            pickle_protocol: specifies pickle protocol when saving, with `torch.save`.
                 Defaults to torch.serialization.DEFAULT_PROTOCOL. For more details, please check:
                 https://pytorch.org/docs/stable/generated/torch.save.html#torch.save.
             hash_transform: a callable to compute hash from the transform information when caching.
@@ -557,7 +557,7 @@ def __init__(
                 defaults to `monai.data.utils.pickle_hashing`.
             db_name: lmdb database file name. Defaults to "monai_cache".
             progress: whether to display a progress bar.
-            pickle_protocol: specifies pickle protocol when saving, with `torch.save`. 
+            pickle_protocol: specifies pickle protocol when saving, with `torch.save`.
                 Defaults to torch.serialization.DEFAULT_PROTOCOL. For more details, please check:
                 https://pytorch.org/docs/stable/generated/torch.save.html#torch.save.
             hash_transform: a callable to compute hash from the transform information when caching.
@@ -601,15 +601,14 @@ def set_data(self, data: Sequence):
         super().set_data(data=data)
         self._read_env = self._fill_cache_start_reader(show_progress=self.progress)
 
-    def _safe_serialize(self,val):
-        out=BytesIO()
-        torch.save(convert_to_tensor(val), out, pickle_protocol =self.pickle_protocol)
+    def _safe_serialize(self, val):
+        out = BytesIO()
+        torch.save(convert_to_tensor(val), out, pickle_protocol=self.pickle_protocol)
         out.seek(0)
         return out.read()
 
-    def _safe_deserialize(self,val):
-        out=BytesIO(val)
-        return torch.load(out,weights_only=True)
+    def _safe_deserialize(self, val):
+        return torch.load(BytesIO(val), map_location="cpu", weights_only=True)
 
     def _fill_cache_start_reader(self, show_progress=True):
         """
@@ -637,7 +636,7 @@ def _fill_cache_start_reader(self, show_progress=True):
                         if val is None:
                             val = self._pre_transform(deepcopy(item))  # keep the original hashed
                             # val = pickle.dumps(val, protocol=self.pickle_protocol)
-                            val=self._safe_serialize(val)
+                            val = self._safe_serialize(val)
                         with env.begin(write=True) as txn:
                             txn.put(key, val)
                         done = True

monai/data/utils.py

@@ -500,7 +500,7 @@ def list_data_collate(batch: Sequence):
     collate_fn = default_collate
     try:
         # if config.USE_META_DICT:
-            # data = pickle_operations(data)  # bc 0.9.0
+        # data = pickle_operations(data)  # bc 0.9.0
         if isinstance(elem, Mapping):
             ret = {}
             for k in elem:
@@ -654,14 +654,14 @@ def decollate_batch(batch, detach: bool = True, pad=True, fill_value=None):
         _gen = zip_longest(*deco.values(), fillvalue=fill_value) if pad else zip(*deco.values())
         ret = [dict(zip(deco, item)) for item in _gen]
         # if not config.USE_META_DICT:
-            # return ret
+        # return ret
         # return pickle_operations(ret, is_encode=False)  # bc 0.9.0
         return ret
     if isinstance(deco, Iterable):
         _gen = zip_longest(*deco, fillvalue=fill_value) if pad else zip(*deco)
         ret_list = [list(item) for item in _gen]
         # if not config.USE_META_DICT:
-            # return ret_list
+        # return ret_list
         # return pickle_operations(ret_list, is_encode=False)  # bc 0.9.0
         return ret_list
     raise NotImplementedError(f"Unable to de-collate: {batch}, type: {type(batch)}.")

monai/utils/state_cacher.py

@@ -64,7 +64,7 @@ def __init__(
             pickle_module: module used for pickling metadata and objects, default to `pickle`.
                 this arg is used by `torch.save`, for more details, please check:
                 https://pytorch.org/docs/stable/generated/torch.save.html#torch.save.
-            pickle_protocol: specifies pickle protocol when saving, with `torch.save`. 
+            pickle_protocol: specifies pickle protocol when saving, with `torch.save`.
                 Defaults to torch.serialization.DEFAULT_PROTOCOL. For more details, please check:
                 https://pytorch.org/docs/stable/generated/torch.save.html#torch.save.
 

monai/utils/type_conversion.py

@@ -117,7 +117,7 @@ def convert_to_tensor(
     wrap_sequence: bool = False,
     track_meta: bool = False,
     safe: bool = False,
-    convert_numeric: bool = True
+    convert_numeric: bool = True,
 ) -> Any:
     """
     Utility to convert the input data to a PyTorch Tensor, if `track_meta` is True, the output will be a `MetaTensor`,