已合并
[Feature] transfer_to_npu adapt cuda_default_generators #34209
louyujing创建于 4月23日
[Feature] transfer_to_npu adapt cuda_default_generators #34209
已合并
共 2 个文件变更+10-0
| @@ -399,6 +399,11 @@ class TestTransferToNpu(TestCase): | |||
| 399 | 399 | ||
| 400 | def test_host_empty_cache_is_patched(self): | 400 | def test_host_empty_cache_is_patched(self): |
| 401 | self.assertEqual(torch._C._host_emptyCache, torch_npu._C._npu_hostEmptyCache) | 401 | self.assertEqual(torch._C._host_emptyCache, torch_npu._C._npu_hostEmptyCache) |
| 402 | + | ||
| 403 | + def test_update_cuda_default_generators(self): | ||
| 404 | + torch.randn(1).npu() | ||
| 405 | + self.assertEqual(torch.cuda.default_generators, torch_npu.npu.default_generators) | ||
| 406 | + self.assertNotEqual(torch.cuda.default_generators, ()) | ||
| 402 | 407 | ||
| 403 | if __name__ == "__main__": | 408 | if __name__ == "__main__": |
| 404 | run_tests() | 409 | run_tests() |
| @@ -423,6 +423,11 @@ def _init(): | |||
| 423 | _device_wrapper(torch.cuda, torch_cuda_fn_white_list) | 423 | _device_wrapper(torch.cuda, torch_cuda_fn_white_list) |
| 424 | torch.cuda.device.__init__ = _wrapper_cuda(torch.cuda.device.__init__) | 424 | torch.cuda.device.__init__ = _wrapper_cuda(torch.cuda.device.__init__) |
| 425 | torch.cuda.amp.autocast_mode = torch_npu.npu.amp.autocast_mode | 425 | torch.cuda.amp.autocast_mode = torch_npu.npu.amp.autocast_mode |
| 426 | + | ||
| 427 | + def _update_cuda_default_generators(): | ||
| 428 | + torch.cuda.default_generators = torch_npu.npu.default_generators | ||
| 429 | + | ||
| 430 | + torch_npu.npu._lazy_call(_update_cuda_default_generators) | ||
| 426 | 431 | ||
| 427 | # torch.cuda.memory.* | 432 | # torch.cuda.memory.* |
| 428 | _device_wrapper(torch.npu.memory, ['_record_memory_history', '_snapshot']) | 433 | _device_wrapper(torch.npu.memory, ['_record_memory_history', '_snapshot']) |