tal: 1min 34s
Wall time: 1min 15s


	Copied
images.shape


	Copied
(8, 1, 512, 512, 3)
We use block_until_ready() to correctly measure inference time, because JAX uses asynchronous dispatch and returns control to the Python loop as soon as it can. You don’t need to use that in your code; blocking will occur automatically when you want to use the result of a computation that has not yet been materialized.
