mxnet.context.gpu¶

mxnet.context.gpu(device_id=0)[source]¶

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.

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)

Parameters: device_id (int, optional) – The device id of the device, needed for GPU.
Returns: context – The corresponding GPU context.
Return type: Context