def __del__(self): with cudart.device(self.device_id): cudart.cuda_stream_destroy(self.cuda_stream)
def __init__(self, device_id=None): with cudart.device(device_id): self.device_id = cudart.cuda_get_device() self.cuda_stream = cudart.ct_cuda_stream() cudart.cuda_stream_create(self.cuda_stream)
def _create_cudnn_handle(device_id): with cudart.device(device_id): handle = cudnn.ct_cudnn_handle() cudnn.create(handle) return handle