已合并
[test] pin_memory support PCIE_through HVA UT #30722
gleaming-spark创建于 2月10日
[test] pin_memory support PCIE_through HVA UT #30722
已合并
共 1 个文件变更+19-0
| @@ -0,0 +1,19 @@ | |||
| 1 | +import unittest | ||
| 2 | + | ||
| 3 | +import os | ||
| 4 | +import torch | ||
| 5 | +from torch.testing._internal.common_utils import TestCase, run_tests | ||
| 6 | +from torch_npu.npu.utils import get_cann_version | ||
| 7 | + | ||
| 8 | + | ||
| 9 | +class TestPinMemoryHostRegister(TestCase): | ||
| 10 | + | ||
| 11 | + | ||
| 12 | + def test_pin_memory_host_register(self): | ||
| 13 | + os.environ["PYTORCH_NPU_ALLOC_CONF"] = "pinned_mem_register:True" | ||
| 14 | + cpu_tensor = torch.ones([2, 3]) | ||
| 15 | + pin_tensor = cpu_tensor.pin_memory() | ||
| 16 | + self.assertTrue(pin_tensor.is_pinned()) | ||
| 17 | + | ||
| 18 | +if __name__ == '__main__': | ||
| 19 | + run_tests() | ||