We read every piece of feedback, and take your input very seriously.
To see all available qualifiers, see our documentation.
There was an error while loading. Please reload this page.
1 parent 18f9bb5 commit 6830aceCopy full SHA for 6830ace
1 file changed
cuda_core/examples/pytorch_example.py
@@ -32,24 +32,14 @@
32
"""
33
34
35
-# Create a wrapper class that implements __cuda_stream__
36
-class PyTorchStreamWrapper:
37
- def __init__(self, pt_stream):
38
- self.pt_stream = pt_stream
39
-
40
- def __cuda_stream__(self):
41
- stream_id = self.pt_stream.cuda_stream
42
- return (0, stream_id) # Return format required by CUDA Python
43
44
45
def main():
46
dev = Device()
47
dev.set_current()
48
49
pt_stream = torch.cuda.current_stream()
50
print(f"PyTorch stream: {pt_stream}", file=sys.stderr)
51
52
- stream = dev.create_stream(PyTorchStreamWrapper(pt_stream))
+ stream = dev.create_stream(pt_stream)
53
54
try:
55
# prepare program
0 commit comments