Skip to content

Commit 6830ace

Browse files
committed
cuda.core: update PyTorch example to use native __cuda_stream__
Signed-off-by: Ralf Juengling <rjuengling@nvidia.com>
1 parent 18f9bb5 commit 6830ace

1 file changed

Lines changed: 1 addition & 11 deletions

File tree

cuda_core/examples/pytorch_example.py

Lines changed: 1 addition & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -32,24 +32,14 @@
3232
"""
3333

3434

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-
4535
def main():
4636
dev = Device()
4737
dev.set_current()
4838

4939
pt_stream = torch.cuda.current_stream()
5040
print(f"PyTorch stream: {pt_stream}", file=sys.stderr)
5141

52-
stream = dev.create_stream(PyTorchStreamWrapper(pt_stream))
42+
stream = dev.create_stream(pt_stream)
5343

5444
try:
5545
# prepare program

0 commit comments

Comments
 (0)