Contrib

Extension Context

nnabla.contrib.context.extension_context(extension_name='cpu', **kw)[source]

Get the context of the specified extension.

All extension’s module must provide context(**kw) function.

Parameters:
  • extension_name (str) – Module path relative to nnabla_ext.

  • kw (dict) – Additional keyword arguments for context function in a extension module.

Returns:

The current extension context.

Return type:

nnabla.Context

Example

ctx = extension_context('cuda.cudnn', device_id=0)
nn.set_default_context(ctx)