Bring Colab notebook up to date with changes in metrics:   - Remove LogSoftmax from end of model.   - Take advantage of non-weighted metrics to simplify data pipelines (since     weights not needed for this model/data).   - Replace CrossEntropyLoss with CategoryCrossEntropy.   - Replace Accuracy with CategoryAccuracy.
Bring Colab notebook up to date with changes in metrics:
  - Remove LogSoftmax from end of model.
  - Take advantage of non-weighted metrics to simplify data pipelines (since
    weights not needed for this model/data).
  - Replace CrossEntropyLoss with CategoryCrossEntropy.
  - Replace Accuracy with CategoryAccuracy.

PiperOrigin-RevId: 346195607

google/trax,google/trax