GreenContext

class torch.cuda.green_contexts.GreenContext

Wrapper around a CUDA green context.

Warning

This API is in beta and may change in future releases.

static create(num_sms, device_id=0)

Create a CUDA green context.

Parameters

• num_sms (int) – The number of SMs to use in the green context.

• device_id (int, optional) – The device index of green context.

Return type

object

pop_context()

Assuming the green context is the current context, pop it from the context stack and restore the previous context.

set_context()

Make the green context the current context.