Use THSetNumThreads instead of omp_set_num_threads

Set OMP num threads to one in the data loader.

Fixes #81
Fixes #82

torch/csrc/Module.cpp

@@ -141,24 +141,15 @@ static PyObject * THPModule_initExtension(PyObject *self, PyObject *shm_manager_
 
 static PyObject * THPModule_getNumThreads(PyObject *module)
 {
-#ifdef _OPENMP
-  return PyLong_FromLong(omp_get_max_threads());
-#else
-  return PyLong_FromLong(1);
-#endif
+  return PyLong_FromLong(THGetNumThreads());
 }
 
 static PyObject * THPModule_setNumThreads(PyObject *module, PyObject *arg)
 {
   THPUtils_assert(THPUtils_checkLong(arg), "set_num_threads expects an int, "
           "but got %s", THPUtils_typename(arg));
-#ifdef _OPENMP
-  omp_set_num_threads(THPUtils_unpackLong(arg));
-#else
-  PyErr_WarnEx(PyExc_RuntimeWarning, "set_num_threads is a no-op - torch was "
-          "compiled without OpenMP support", 1);
-#endif
-  return 0;
+  THSetNumThreads((int)THPUtils_unpackLong(arg));
+  Py_RETURN_NONE;
 }
 
 bool THPModule_isTensor(PyObject *obj)

torch/utils/data/dataloader.py

@@ -21,6 +21,7 @@ def _processBatch(dataset, indices, collate_fn):
 
 
 def _workerLoop(dataset, index_queue, data_queue, collate_fn):
+    torch.set_num_threads(1)
     while True:
         batch_indices = index_queue.get()