Fanli LinandCopilot
bfbf3e5b69
Signed-off-by: Fanli Lin <fanli.lin@intel.com> Signed-off-by: Fanli Lin <fanli0116@gmail.com> Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
17 lines
428 B
Python
17 lines
428 B
Python
from contextlib import contextmanager
|
|
|
|
import torch
|
|
|
|
|
|
@contextmanager
|
|
def torch_cuda_wrapper():
|
|
try:
|
|
# replace cuda APIs with xpu APIs, this should work by default
|
|
torch.cuda.Stream = torch.xpu.Stream
|
|
torch.cuda.default_stream = torch.xpu.current_stream
|
|
torch.cuda.current_stream = torch.xpu.current_stream
|
|
torch.cuda.stream = torch.xpu.stream
|
|
yield
|
|
finally:
|
|
pass
|