Rate this Page

GreenContext#

class torch.cuda.green_contexts.GreenContext[source]#

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)[source]#

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

make_current()[source]#

Make the green context the current context.

pop_current()[source]#

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