Update PyTorch example to skip using a custom stream adaptor

The upstream has implemented __cuda_stream__ support 🎉
pytorch/pytorch#163614
Once a new release of PyTorch is out, we can update our example and check if the custom adapter is still needed, or if we can use the built-in support:

cuda-python/cuda_core/examples/pytorch_example.py

Lines 40 to 50 in 279c943

    
           # 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 
        
           s = dev.create_stream(PyTorchStreamWrapper(pt_stream))

Navigation Menu

Repository navigation

Update PyTorch example to skip using a custom stream adaptor #1083

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

	# 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


	s = dev.create_stream(PyTorchStreamWrapper(pt_stream))

Update PyTorch example to skip using a custom stream adaptor #1083

Description

Activity

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions