Getting a shape mismatch for the EIM Loss

I'm trying to run a model with the EIM loss but i get an index error: shape mismatch for the line 

      probabilities_at_targets=probabilities_for_training[range(distances.size(0)), targets] 

I don't know if anyone has faced this problem before and/or knows how to solve it?

Thanks a lot in advance,

Guillem