[Task] Figure out why the unit test for implicit fails on `np.allclose` (it passes in the notebook when run on real data)

### Description

`np.allclose` fails in the unit test when comparing local predictions vs predictions received from triton. It doesn't fail when implicit is trained on real (and non-synthetic) data.

This task is to figure out what is going on there.