AttributeError: module 'dill._dill' has no attribute 'stack'

[yangjenhao]

NVIDIA-SMI 460.32.03 Driver Version: 460.32.03 CUDA Version: 11.2
using your notebooks google colab
https://colab.research.google.com/drive/12LjJazBl7Gam0XBPy_y0CTOJZeZ34c2v?usp=sharing
when do this

train_dataset = train_dataset.map(
    process_data_to_model_inputs,
    batched=True,
    batch_size=batch_size,
    remove_columns=["article", "abstract", "section_names"],
)

---

AttributeError Traceback (most recent call last)
in
3 batched=True,
4 batch_size=batch_size,
----> 5 remove_columns=["article", "abstract", "section_names"],
6 )

12 frames
/usr/local/lib/python3.7/dist-packages/datasets/utils/py_utils.py in save_function(pickler, obj)
529 globalvars = dill.detect.globalvars
530 globs = globalvars(obj, recurse=True, builtin=True)
--> 531 if id(obj) in dill._dill.stack:
532 globs = obj.globals if dill._dill.PY3 else obj.func_globals
533 else:

AttributeError: module 'dill._dill' has no attribute 'stack'

---

And I already try this
pip install dill==0.3.4
but not work