在TensorFlow的Eager模式下阻止GPU操作的最佳方法是什么？

Question

我想知道在TensorFlow Eager模式下等待GPU操作完成的推荐方法。

位于GPU设备上的操作似乎是异步执行的（我在TensorFlow文档中找不到此操作，但这与行为一致）。例如，在使用time.time() *对GPU操作进行计时时，这一点很重要，因为我们需要确保在记录结束时间之前操作已完成。

我能找到确保已执行GPU操作的唯一方法是将输出数据（部分）显式复制到CPU。

例如（假设所有操作均在GPU上执行）：

t0 = time.time()
result = f(input_tensor)  # carry out some operations on the input
_ = result[0].numpy()  # copies a single element of the output tensor to the CPU
t1 = time.time()
print("runtime =", t1 - t0)

由于将数据复制到CPU会产生一些开销，因此最好有一种方法来确保GPU在不复制的情况下完成执行。有这种方法吗？也许像JAX的block_until_ready()一样？

*我意识到使用time.time()可能不是在Eager模式下计时GPU操作的最佳方法。