Setting NPU as Default Device

The TF Adapter provides an API for registering an NPU as a valid TensroFlow device. You can add the following code to the beginning of the script file that is decorated with @tf.function and can work properly under the CPU or GPU to set an NPU as the default device:

1
2
3

import npu_device as npu
# By default, device 0 is a compute device.
npu.open().as_default()

Perform this step before importing other Python modules. Otherwise, operators might be scheduled to devices other than your NPU.

For details about npu.open, see npu.open.