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

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)
[](https://localhost:8080/#) 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](https://localhost:8080/#) 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