Plot mismatch when using trained model for prediction

[marcemq]

Issue

When executing train_torch_deterministic with --id-test=0 at prediction step, the resulting plot showcase a mismatch between the reference image and the predicted path, as shown below:

When using --id-test=3 you can't tell the difference since the reference image is almost the same

Steps to reproduce

At repository directory execute:

➜  trajpred-bdl git:(main) python3 tests/train_torch_deterministic.py --batch-size=64 --epochs=10 --examples=2 --pickle --id-test=0

Expected result

The predicted path should be drown in the reference image.