Loss Fn - cosine similarity

[cksac]

ref: https://optax.readthedocs.io/en/latest/api/losses.html#optax.cosine_similarity