diff --git a/tests/L0/device.py b/tests/L0/device.py index 5f103cc0413d195559f4f34e4d45352e73a82dab..1b8e05ad1c6114aa621255e464e08c27851a0a30 100644 --- a/tests/L0/device.py +++ b/tests/L0/device.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. + import torch if torch.cuda.is_available(): CALCULATE_DEVICE = 'cuda' @@ -19,4 +20,4 @@ else: CALCULATE_DEVICE = 'npu' def is_npu(): - return CALCULATE_DEVICE.find('npu') != -1 \ No newline at end of file + return CALCULATE_DEVICE.find('npu') != -1