Better CMS dataset, fix f16, fix transformer

[jpata]

• switch to dataset 1.3.1, removing PS and BREM (reduces size to ~5k elements / event): jpata/hep_tfds@05ac2ef
• fix type conversions for f16 training. Still gives a nan (square goes out of range for f16), but testable.
• all loss coefs to 1.0, use PCGrad for multi-task optimization
• fix the transformer model so that it can train
• disable per-event energy loss for now (not sure if it's really useful)
• benchmark training with 1.3.1 here: https://www.comet.ml/jpata/particleflow-tf/0f3ad9d2885f42d8bef48891d14e3071