Table Of Contents
Table Of Contents

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