Fix order of shape for unbatchify in POMO

[hyeok9855]

Description

Fix order of argument shape for unbatchify function in POMO Model.

Types of changes

• Bug fix (non-breaking change which fixes an issue)

The original argument for shape was (n_start, n_aug), with which the shape of tensors will be changed to [batch_size, n_start, n_aug]. But we want it to be [batch_size, n_aug, n_start] so that it can mirror the batchify operation appropriately.

Motivation and Context

Note that the bug will not harm your training procedure or maximum reward calculation, but it should be fixed because it is not what we wanted, and make debugging difficult.

Checklist

Go over all the following points, and put an x in all the boxes that apply.
If you are unsure about any of these, don't hesitate to ask. We are here to help!

I found that there's no test code for POMO though :(

• My change requires a change to the documentation.
• I have updated the tests accordingly (required for a bug fix or a new feature).
• I have updated the documentation accordingly.

[codecov]

Codecov Report

Attention: 6 lines in your changes are missing coverage. Please review.

Files	Coverage Δ
rl4co/models/zoo/pomo/model.py	20.40% <0.00%> (-0.43%)	⬇️

📢 Thoughts on this report? Let us know!.

[hyeok9855]

By the way, I guess the selecting best actions are meaningless because it is not used or returned.

[fedebotu]

Thanks!
Yes as we discussed, this part was working anyways since it's only during inference (so it doesn't matter in which order how you take the max(max()). But it's better to fix ofc 👌
Btw could you have a look at SymNCO as well to see if they there is the same problem?