Comment by cgarciae

Side note: I like Pytorch but eager pytorch is not faster the jax.jit or tf.function code