diff --git a/CHANGELOG.md b/CHANGELOG.md index 5a91183ab..77f88f2e6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -25,6 +25,7 @@ - PR #421 Capture thread id in logging and improve logger testing - PR #426 Added multi-threaded support to replay benchmark. - PR #435 Update conda upload versions for new supported CUDA/Python +- PR #437 Test with `pickle5` (for older Python versions) ## Bug Fixes diff --git a/python/rmm/tests/test_rmm.py b/python/rmm/tests/test_rmm.py index 3f4051738..5411ecf45 100644 --- a/python/rmm/tests/test_rmm.py +++ b/python/rmm/tests/test_rmm.py @@ -1,6 +1,5 @@ # Copyright (c) 2020, NVIDIA CORPORATION. -import pickle import sys import tempfile from itertools import product @@ -11,6 +10,14 @@ import rmm +if sys.version_info < (3, 8): + try: + import pickle5 as pickle + except ImportError: + import pickle +else: + import pickle + cuda.set_memory_manager(rmm.RMMNumbaManager)