xla

1f7c1072 - update batch size to 512 because of tracing issue