OSSForks/pytorch
1
0
Fork 0
mirror of https://github.com/zebrajr/pytorch.git synced 2026-01-15 12:15:51 +00:00
Code Issues Actions 52 Packages Projects Releases Wiki Activity

chunky sync - build scripts to be written

Browse Source
This commit is contained in:
Yangqing Jia
2016-07-21 10:16:42 -07:00
parent 292211baa4
commit 6463eebc7b
223 changed files with 12059 additions and 5320 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
caffe2/python/caffe2_python_gpu.cc
View File

@@ -20,15 +20,8 @@ extern "C" {
// Here are functions that are purely GPU-based functions to be filled.
PyObject* NumberOfGPUs(PyObject* self, PyObject* args) {
int num_devices = 0;
auto err = cudaGetDeviceCount(&num_devices);
if (err == cudaErrorNoDevice || err == cudaErrorInsufficientDriver) {
return Py_BuildValue("i", 0);
} else if (err != cudaSuccess) {
PyErr_SetString(PyExc_RuntimeError, "Runtime CUDA error.");
return nullptr;
}
PyObject* NumCudaDevices(PyObject* self, PyObject* args) {
int num_devices = caffe2::NumCudaDevices();
return Py_BuildValue("i", num_devices);
}
@@ -77,7 +70,7 @@ static PyMethodDef* GetCaffe2PythonGPUMethods() {
// TODO(Yangqing): write the methods string.
// Note(Yangqing): For any function that we are going to override in the
// python file, we prepend "cc_" here.
_PYNAME(NumberOfGPUs),
_PYNAME(NumCudaDevices),
_PYNAME(SetDefaultGPUID),
_PYNAME(GetDefaultGPUID),
_PYNAME(GetCudaPeerAccessPattern),