xla

a4237f08 - move the model in backward pass to xla device.