mirror of
https://gitee.com/ascend/ModelLink.git
synced 2024-12-04 04:48:27 +08:00
8 lines
192 B
Python
8 lines
192 B
Python
|
import torch
|
||
|
|
import torch_npu
|
||
|
|
from torch_npu.contrib import transfer_to_npu
|
||
|
|
|
||
|
|
|
||
|
|
torch.cuda.init = torch.npu.init
|
||
|
|
torch.npu.init()
|
||
|
|
torch.cuda.default_generators = torch_npu.npu.default_generators
|