From 289db03e2a3078400dd9f8f46c8b53ac3b5e40fb Mon Sep 17 00:00:00 2001 From: Ralf Juengling Date: Mon, 6 Jul 2026 16:39:10 -0700 Subject: [PATCH 1/2] cuda.core: update PyTorch example to use native __cuda_stream__ Signed-off-by: Ralf Juengling --- cuda_core/examples/pytorch_example.py | 12 +----------- 1 file changed, 1 insertion(+), 11 deletions(-) diff --git a/cuda_core/examples/pytorch_example.py b/cuda_core/examples/pytorch_example.py index 5826c0a4421..d235dc58452 100644 --- a/cuda_core/examples/pytorch_example.py +++ b/cuda_core/examples/pytorch_example.py @@ -32,16 +32,6 @@ """ -# Create a wrapper class that implements __cuda_stream__ -class PyTorchStreamWrapper: - def __init__(self, pt_stream): - self.pt_stream = pt_stream - - def __cuda_stream__(self): - stream_id = self.pt_stream.cuda_stream - return (0, stream_id) # Return format required by CUDA Python - - def main(): dev = Device() dev.set_current() @@ -49,7 +39,7 @@ def main(): pt_stream = torch.cuda.current_stream() print(f"PyTorch stream: {pt_stream}", file=sys.stderr) - stream = dev.create_stream(PyTorchStreamWrapper(pt_stream)) + stream = dev.create_stream(pt_stream) try: # prepare program From a23ff1f9506bbc4f776a793b89264734c36d77f0 Mon Sep 17 00:00:00 2001 From: Ralf Juengling Date: Tue, 7 Jul 2026 16:08:29 -0700 Subject: [PATCH 2/2] [] Add torch version check --- cuda_core/examples/pytorch_example.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/cuda_core/examples/pytorch_example.py b/cuda_core/examples/pytorch_example.py index d235dc58452..480e18dfe98 100644 --- a/cuda_core/examples/pytorch_example.py +++ b/cuda_core/examples/pytorch_example.py @@ -32,10 +32,22 @@ """ +def torch_version() -> tuple[int, int, int]: + v = torch.__version__ + v = v.split("+")[0].split(".") + return tuple(int(x) for x in v) + + def main(): dev = Device() dev.set_current() + print(f"PyTorch version: {torch_version()}", file=sys.stderr) + if torch_version() < (2, 10, 0): + # Need support for the __cuda_stream__ method. + print("PyTorch version 2.10 or later required, skipping example", file=sys.stderr) + return + pt_stream = torch.cuda.current_stream() print(f"PyTorch stream: {pt_stream}", file=sys.stderr)