gpu

gpu(device_id=0)

Returns a GPU context.

This function is a short cut for Context(‘gpu’, device_id). The K GPUs on a node are typically numbered as 0,…,K-1.

Parameters:

device_id : int, optional

The device id of the device, needed for GPU.

Returns:

context : Context

The corresponding GPU context.

Examples

>>> cpu_array = mx.nd.ones((2, 3))
>>> cpu_array.context
cpu(0)
>>> with mx.gpu(1):
...     gpu_array = mx.nd.ones((2, 3))
>>> gpu_array.context
gpu(1)
>>> gpu_array = mx.nd.ones((2, 3), ctx=mx.gpu(1))
>>> gpu_array.context
gpu(1)