Train Model with XLA JIT Enabled

Run the Next Cell to Show XLA JIT Training Code

Keep an eye on the following:

with tf.device('device:XLA_GPU:0')



In [ ]:

    
cat /root/src/main/python/xla/train_linear_xla.py

The cell below will report

Aborted (core dumped)

but still generates useful info.



In [ ]:

    
%%bash

train_xla