设置NPU为默认设备

在带有@tf.function装饰并在CPU或GPU下可以正常工作的脚本文件开头,加载npu_device并设置NPU为默认设备。

import npu_device as npu
npu.open().as_default()

您应当在import其他python包前执行该操作,防止在加载后续包的过程中有未分发到NPU的算子执行行为。