Fixing pickle protocol issue

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

Author: ericspod

tests/data/test_gdsdataset.py

@@ -86,7 +86,8 @@ def test_cache(self):
                     cache_dir=tempdir,
                     device=0,
                     pickle_module="pickle",
-                    pickle_protocol=pickle.HIGHEST_PROTOCOL,
+                    # TODO: was pickle.HIGHEST_PROTOCOL but this wasn't compatible with torch.load, need to improve compatibility
+                    pickle_protocol=torch.serialization.DEFAULT_PROTOCOL,
                 )
                 assert_allclose(items[0], p(np.arange(0, np.prod(shape)).reshape(shape)))
                 ds1 = GDSDataset(items, transform=_InplaceXform(), cache_dir=tempdir, device=0)

tests/data/test_persistentdataset.py

@@ -12,12 +12,12 @@
 from __future__ import annotations
 
 import os
-import pickle
 import tempfile
 import unittest
 
 import nibabel as nib
 import numpy as np
+import torch
 from parameterized import parameterized
 
 from monai.data import PersistentDataset, json_hashing
@@ -66,7 +66,8 @@ def test_cache(self):
                 transform=_InplaceXform(),
                 cache_dir=tempdir,
                 pickle_module="pickle",
-                pickle_protocol=pickle.HIGHEST_PROTOCOL,
+                # TODO: was pickle.HIGHEST_PROTOCOL but this wasn't compatible with torch.load, need to improve compatibility
+                pickle_protocol=torch.serialization.DEFAULT_PROTOCOL,
             )
             self.assertEqual(items, [[[]], [[0]], [[0, 1]], [[0, 1, 2]], [[0, 1, 2, 3]]])
             ds1 = PersistentDataset(items, transform=_InplaceXform(), cache_dir=tempdir)