Tensor shape error in sac_continuous_action_torchcompile.py on Mac

[ran-weii]

Hi, thanks for this cool project!

I was running sac_continuous_action_torchcompile.py and wanted to do a comparison of compile vs no compile on Mac M3. I changed line 172 to use MPS instead of CUDA.

I ran into a weird issue on line 222 which calculates the MSE loss:

RuntimeError: impl: target and input tensors must have identical shapes

I checked val.view(-1) and next_q_value actually have the same shape. So I don't understand why this error happens.

I was able to solve it by changing to the following loss computation:

loss_val = torch.pow(vals.view(-1) - next_q_value, 2).mean()

I don't know if this will affect performance significantly (I doubt).

I installed both torchrl and tensordict from source if that matters.

[vmoens]

I raised an issue with this: pytorch/pytorch#140704
Let's hope it gets fixed!