Python caffe2.python.workspace.NumCudaDevices() Examples

The following are 1 code examples of caffe2.python.workspace.NumCudaDevices(). You can vote up the ones you like or vote down the ones you don't like, and go to the original project or source file by following the links above each example. You may also want to check out all available functions/classes of the module caffe2.python.workspace , or try the search function .
Example #1
Source File: model.py    From dlcookbook-dlbs with Apache License 2.0 6 votes vote down vote up 
def get_device_option(gpu=None):
        """Constructs `core.DeviceOption` object

        :param int gpu: Identifier of GPU to use or None for CPU.
        :return: Instance of `core.DeviceOption`.
        """
        dev_opt = None
        if gpu is None:
            dev_opt = core.DeviceOption(caffe2_pb2.CPU)
        else:
            assert workspace.has_gpu_support, "Workspace does not support GPUs"
            assert gpu >= 0 and gpu < workspace.NumCudaDevices(),\
                   "Workspace does not provide this gpu (%d). "\
                   "Number of GPUs is %d" % (gpu, workspace.NumCudaDevices())
            dev_opt = core.DeviceOption(caffe2_pb2.CUDA, gpu)
        return dev_opt