import torch

import torch_npu

import time


print(torch_npu.npu.is_available())

torch_npu.npu.set_device(0)

start_time = time.perf_counter()

device = torch.device("npu" if torch_npu.npu.is_available() else "cpu")

print("using device:", device)

x = torch.randn(2, 2).npu()

y = torch.randn(2, 2).npu()

z = x.mm(y)

print(z)

end_time = time.perf_counter()

duration_ms = (end_time - start_time) * 1000

print(f"Duration: {duration_ms:.3f} ms")

python npu.py
/home/HwHiAiUser/.local/lib/python3.9/site-packages/torch_npu/dynamo/__init__.py:18: UserWarning: Register eager implementation for the 'npu' backend of dynamo, as torch_npu was not compiled with torchair.
  warnings.warn(
True
using device: npu
Warning: Device do not support double dtype now, dtype cast repalce with float.
tensor([[-0.0325,  1.3582],
        [-1.0262, -1.6385]], device='npu:0')
Duration: 11189.677 ms