Hi, I'm getting this error when using python -m osu_dreamer predict

[Zheyuan-Yang]

Hi, I'm getting this error when using python -m osu_dreamer predict

Traceback (most recent call last):
  File "/home/andrewyang/ssd_ubuntu/anaconda3/envs/osu-dreamer/lib/python3.9/runpy.py", line 197, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/home/andrewyang/ssd_ubuntu/anaconda3/envs/osu-dreamer/lib/python3.9/runpy.py", line 87, in _run_code
    exec(code, run_globals)
  File "/home/andrewyang/ssd_ubuntu/osu-dreamer/osu_dreamer/__main__.py", line 20, in <module>
    main()
  File "/home/andrewyang/ssd_ubuntu/anaconda3/envs/osu-dreamer/lib/python3.9/site-packages/click/core.py", line 1157, in __call__
    return self.main(*args, **kwargs)
  File "/home/andrewyang/ssd_ubuntu/anaconda3/envs/osu-dreamer/lib/python3.9/site-packages/click/core.py", line 1078, in main
    rv = self.invoke(ctx)
  File "/home/andrewyang/ssd_ubuntu/anaconda3/envs/osu-dreamer/lib/python3.9/site-packages/click/core.py", line 1688, in invoke
    return _process_result(sub_ctx.command.invoke(sub_ctx))
  File "/home/andrewyang/ssd_ubuntu/anaconda3/envs/osu-dreamer/lib/python3.9/site-packages/click/core.py", line 1434, in invoke
    return ctx.invoke(self.callback, **ctx.params)
  File "/home/andrewyang/ssd_ubuntu/anaconda3/envs/osu-dreamer/lib/python3.9/site-packages/click/core.py", line 783, in invoke
    return __callback(*args, **kwargs)
  File "/home/andrewyang/ssd_ubuntu/anaconda3/envs/osu-dreamer/lib/python3.9/site-packages/jaxtyping/_decorator.py", line 453, in wrapped_fn
    out = fn(*args, **kwargs)
  File "/home/andrewyang/ssd_ubuntu/osu-dreamer/osu_dreamer/scripts/predict.py", line 56, in predict
    model = Model.load_from_checkpoint(model_path).eval()
  File "/home/andrewyang/ssd_ubuntu/anaconda3/envs/osu-dreamer/lib/python3.9/site-packages/pytorch_lightning/utilities/model_helpers.py", line 125, in wrapper
    return self.method(cls, *args, **kwargs)
  File "/home/andrewyang/ssd_ubuntu/anaconda3/envs/osu-dreamer/lib/python3.9/site-packages/pytorch_lightning/core/module.py", line 1582, in load_from_checkpoint
    loaded = _load_from_checkpoint(
  File "/home/andrewyang/ssd_ubuntu/anaconda3/envs/osu-dreamer/lib/python3.9/site-packages/pytorch_lightning/core/saving.py", line 91, in _load_from_checkpoint
    model = _load_state(cls, checkpoint, strict=strict, **kwargs)
  File "/home/andrewyang/ssd_ubuntu/anaconda3/envs/osu-dreamer/lib/python3.9/site-packages/pytorch_lightning/core/saving.py", line 165, in _load_state
    obj = instantiator(cls, _cls_kwargs) if instantiator else cls(**_cls_kwargs)
  File "/home/andrewyang/ssd_ubuntu/anaconda3/envs/osu-dreamer/lib/python3.9/site-packages/jaxtyping/_decorator.py", line 413, in wrapped_fn
    bound = param_signature.bind(*args, **kwargs)
  File "/home/andrewyang/ssd_ubuntu/anaconda3/envs/osu-dreamer/lib/python3.9/inspect.py", line 3045, in bind
    return self._bind(args, kwargs)
  File "/home/andrewyang/ssd_ubuntu/anaconda3/envs/osu-dreamer/lib/python3.9/inspect.py", line 2960, in _bind
    raise TypeError(msg) from None
TypeError: missing a required argument: 'val_batches'

[Zheyuan-Yang]

Is there anything missing in the checkpoint?