Can't pickle local object 'to_map_style_dataset.<locals>._MapStyleDataset'

🐛 Bug

Describe the bug Can't pickle local object 'to_map_style_dataset.<locals>._MapStyleDataset' error is produced while saving the entire model with torch.save(model, "model.pth")

To Reproduce Steps to reproduce the behavior:

1. Download the script - https://github.com/mlflow/mlflow-torchserve/blob/0e93a3e41b5a4ba1b8cf501a95f681cfb6cea24d/examples/BertNewsClassification/news_classifier.py . The script was adapted from the pytorch examples
2. Install dependent libraries using pip (torch, torchtext, mlflow, transformers, numpy)
3. Run the script with following command - pip install news_classifier.py --max_epochs 1 --num_samples 100
4. The script uses mlflow.pytorch library to save the model. mlflow.pytorch internally saves the entire model using torch.save. The following exception is thrown

AttributeError: Can't pickle local object 'to_map_style_dataset.<locals>._MapStyleDataset'

Looks like the to_map_style_dataset is not serializable.

Expected behavior Model should be saved without any error

Screenshots If applicable, add screenshots to help explain your problem.

Environment

Collecting environment information...
PyTorch version: 1.12.0+cu102
Is debug build: False
CUDA used to build PyTorch: 10.2
ROCM used to build PyTorch: N/A

OS: Ubuntu 16.04.7 LTS (x86_64)
GCC version: (Ubuntu 5.4.0-6ubuntu1~16.04.12) 5.4.0 20160609
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.23

Python version: 3.8.2 (default, Mar 26 2020, 15:53:00)  [GCC 7.3.0] (64-bit runtime)
Python platform: Linux-4.15.0-142-generic-x86_64-with-glibc2.10
Is CUDA available: False
CUDA runtime version: No CUDA
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] botorch==0.5.0
[pip3] gpytorch==1.5.0
[pip3] mlflow-torchserve==0.2.0
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.23.1
[pip3] numpydoc==1.1.0
[pip3] pytorch-kfp-components==4.3.15
[pip3] pytorch-lightning==1.6.5
[pip3] pytorch-ranger==0.1.1
[pip3] torch==1.12.0
[pip3] torch-model-archiver==0.6.0
[pip3] torch-optimizer==0.1.0
[pip3] torch-tb-profiler==0.2.1
[pip3] torch-workflow-archiver==0.2.2b20220224
[pip3] torchaudio==0.10.0+cpu
[pip3] torchdata==0.4.0
[pip3] torchmetrics==0.7.3
[pip3] torchserve==0.6.0
[pip3] torchtext==0.13.0
[pip3] torchvision==0.13.0
[pip3] vit-pytorch==0.35.2
[conda] _tflow_select             2.3.0                       mkl  
[conda] blas                      1.0                         mkl  
[conda] botorch                   0.5.0                    pypi_0    pypi
[conda] cudatoolkit               10.2.89              hfd86e86_1  
[conda] gpytorch                  1.5.0                    pypi_0    pypi
[conda] mkl                       2020.2                      256  
[conda] mkl-service               2.3.0            py38he904b0f_0  
[conda] mkl_fft                   1.2.0            py38h23d657b_0  
[conda] mkl_random                1.1.1            py38h0573a6f_0  
[conda] mlflow-torchserve         0.2.0                    pypi_0    pypi
[conda] numpy                     1.23.1                   pypi_0    pypi
[conda] numpy-base                1.19.2           py38hfa32c7d_0  
[conda] numpydoc                  1.1.0                      py_0  
[conda] pytorch-kfp-components    4.3.15                   pypi_0    pypi
[conda] pytorch-lightning         1.6.5                    pypi_0    pypi
[conda] pytorch-ranger            0.1.1                    pypi_0    pypi
[conda] tensorflow-base           2.2.0           mkl_py38h5059a2d_0  
[conda] torch                     1.12.0                   pypi_0    pypi
[conda] torch-model-archiver      0.6.0                    pypi_0    pypi
[conda] torch-optimizer           0.1.0                    pypi_0    pypi
[conda] torch-tb-profiler         0.2.1                    pypi_0    pypi
[conda] torch-workflow-archiver   0.2.2b20220224           pypi_0    pypi
[conda] torchaudio                0.10.0+cpu               pypi_0    pypi
[conda] torchdata                 0.4.0                    pypi_0    pypi
[conda] torchmetrics              0.7.3                    pypi_0    pypi
[conda] torchserve                0.6.0                    pypi_0    pypi
[conda] torchtext                 0.13.0                   pypi_0    pypi
[conda] torchvision               0.13.0                   pypi_0    pypi
[conda] vit-pytorch               0.35.2                   pypi_0    pypi

Additional context Add any other context about the problem here.

Sample Logs - https://gist.github.com/shrinath-suresh/00086bf503690dd0365e39846ef2dbb5

[parmeet]

I moved the prepare data out of model class and the issue is resolved.

Sounds good @shrinath-suresh, let me close this issue for now then.

[shrinath-suresh]

@parmeet @ejguan Thanks for your input. I moved the prepare data out of model class and the issue is resolved. code link here