[BUG] CVRPlib inference bug.

[Elon-Lau]

I use the same code in https://github.com/ai4co/rl4co/blob/main/examples/2d-meta_train.py for training and the same code in https://github.com/ai4co/rl4co/blob/main/examples/datasets/2-test-on-cvrplib.ipynb for cvrplib instance inference. But the following bug still occur.

Traceback (most recent call last): File "/NCO/RL4CO/examples/test_cvrplib_demo.py", line 91, in out = policy(td.clone(), decode_type='greedy', return_actions=True) File "/miniconda3/envs/rl4co/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl return self._call_impl(*args, *kwargs) File "/miniconda3/envs/rl4co/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl return forward_call(args, **kwargs) File "/miniconda3/envs/rl4co/lib/python3.9/site-packages/rl4co/models/common/constructive/base.py", line 237, in forward td = env.step(td)["next"] File "/miniconda3/envs/rl4co/lib/python3.9/site-packages/rl4co/envs/common/base.py", line 129, in step td = self._step(td) File "/miniconda3/envs/rl4co/lib/python3.9/site-packages/rl4co/envs/routing/cvrp/env.py", line 89, in _step visited = td["visited"].scatter(-1, current_node, 1) RuntimeError: Index tensor must have the same number of dimensions as self tensor

[Elon-Lau]

@fedebotu Can you help me look at the problem, please. Many thanks

[fedebotu]

Thanks for reporting the bug @Elon-Lau - you may check the most recent improvements! Please let us know if now it works well on your end ~

[Elon-Lau]

Thanks for the update, it looks like the patch will fix the current bug.