Setting NPU as Default Device

The TF Adapter provides an API for registering an NPU as a valid TensorFlow device. You can add the following code to the beginning of the script file that has the @tf.function decorator 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 operation before importing other Python packages. Otherwise, operators may be executed on devices other than the NPU during the loading of subsequent packages.

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