xla

b401b2a9 - Add simple Transformer fine-tuning test (#4593)

Add simple Transformer fine-tuning test (#4593) * Add simple Transformer fine-tuning test * Clean up * formatting * Remove extra import Co-authored-by: Doruk Arisoy <darisoy@google.com>