meshgraphnets: ragged rank error when training flag_dynamic

python: conda env 3.7
command:
```
python -m meshgraphnets.run_model --mode=train --model=cloth --checkpoint_dir=./meshgraphnets/chk/dynamic --dataset_dir=./meshgraphnets/data/flag_dynamic
```

pip list output:
```
Package Version
---------------------- ---------
absl-py 1.0.0
astor 0.8.1
certifi 2021.10.8
cloudpickle 1.1.1
contextlib2 21.6.0
cycler 0.11.0
decorator 5.1.1
dm-sonnet 1.36
fonttools 4.29.0
gast 0.2.2
google-pasta 0.2.0
grpcio 1.43.0
h5py 2.10.0
importlib-metadata 4.10.1
Keras-Applications 1.0.8
Keras-Preprocessing 1.1.2
kiwisolver 1.3.2
Markdown 3.3.6
matplotlib 3.5.1
numpy 1.18.5
opt-einsum 3.3.0
packaging 21.3
Pillow 9.0.0
pip 21.2.4
protobuf 3.19.3
pyparsing 3.0.7
python-dateutil 2.8.2
semantic-version 2.8.5
setuptools 58.0.4
six 1.16.0
tensorboard 1.15.0
tensorflow-estimator 1.15.1
tensorflow-gpu 1.15.5
tensorflow-probability 0.8.0
termcolor 1.1.0
typing_extensions 4.0.1
Werkzeug 2.0.2
wheel 0.37.1
wincertstore 0.2
wrapt 1.13.3
zipp 3.7.0
```

error message:
```
Traceback (most recent call last):
File "D:\Anaconda3\envs\mgn\lib\runpy.py", line 193, in _run_module_as_main
"__main__", mod_spec)
File "D:\Anaconda3\envs\mgn\lib\runpy.py", line 85, in _run_code
exec(code, run_globals)
File "E:\Downloads\meshgraphnets\run_model.py", line 131, in
app.run(main)
File "D:\Anaconda3\envs\mgn\lib\site-packages\absl\app.py", line 312, in run
_run_main(main, args)
File "D:\Anaconda3\envs\mgn\lib\site-packages\absl\app.py", line 258, in _run_main
sys.exit(main(argv))
File "E:\Downloads\meshgraphnets\run_model.py", line 126, in main
learner(model, params)
File "E:\Downloads\meshgraphnets\run_model.py", line 61, in learner
noise_gamma=params['gamma'])
File "E:\Downloads\meshgraphnets\dataset.py", line 86, in split_and_preprocess
ds = ds.map(add_noise, num_parallel_calls=8)
File "D:\Anaconda3\envs\mgn\lib\site-packages\tensorflow_core\python\data\ops\dataset_ops.py", line 1913, in map
self, map_func, num_parallel_calls, preserve_cardinality=False))
File "D:\Anaconda3\envs\mgn\lib\site-packages\tensorflow_core\python\data\ops\dataset_ops.py", line 3472, in __init__
use_legacy_function=use_legacy_function)
File "D:\Anaconda3\envs\mgn\lib\site-packages\tensorflow_core\python\data\ops\dataset_ops.py", line 2713, in __init__
self._function = wrapper_fn._get_concrete_function_internal()
File "D:\Anaconda3\envs\mgn\lib\site-packages\tensorflow_core\python\eager\function.py", line 1853, in _get_concrete_function_internal
*args, **kwargs)
File "D:\Anaconda3\envs\mgn\lib\site-packages\tensorflow_core\python\eager\function.py", line 1847, in _get_concrete_function_internal_garbage_collected
graph_function, _, _ = self._maybe_define_function(args, kwargs)
File "D:\Anaconda3\envs\mgn\lib\site-packages\tensorflow_core\python\eager\function.py", line 2147, in _maybe_define_function
graph_function = self._create_graph_function(args, kwargs)
File "D:\Anaconda3\envs\mgn\lib\site-packages\tensorflow_core\python\eager\function.py", line 2038, in _create_graph_function
capture_by_value=self._capture_by_value),
File "D:\Anaconda3\envs\mgn\lib\site-packages\tensorflow_core\python\framework\func_graph.py", line 915, in func_graph_from_py_func
func_outputs = python_func(*func_args, **func_kwargs)
File "D:\Anaconda3\envs\mgn\lib\site-packages\tensorflow_core\python\data\ops\dataset_ops.py", line 2707, in wrapper_fn
ret = _wrapper_helper(*args)
File "D:\Anaconda3\envs\mgn\lib\site-packages\tensorflow_core\python\data\ops\dataset_ops.py", line 2648, in _wrapper_helper
self._input_structure, args)
File "D:\Anaconda3\envs\mgn\lib\site-packages\tensorflow_core\python\data\util\structure.py", line 226, in from_compatible_tensor_list
element_spec, tensor_list)
File "D:\Anaconda3\envs\mgn\lib\site-packages\tensorflow_core\python\data\util\structure.py", line 201, in _from_tensor_list_helper
flat_ret.append(decode_fn(component_spec, value))
File "D:\Anaconda3\envs\mgn\lib\site-packages\tensorflow_core\python\data\util\structure.py", line 225, in
lambda spec, value: spec._from_compatible_tensor_list(value),
File "D:\Anaconda3\envs\mgn\lib\site-packages\tensorflow_core\python\ops\ragged\ragged_tensor.py", line 2027, in _from_compatible_tensor_list
"ragged_rank must be non-negative; got %s." % self._ragged_rank)
ValueError: ragged_rank must be non-negative; got 0.
```
Dataset flag_minimal and sphere_minimal works.
I've tried the same config on my Ubuntu machine but got the same error.