misc update
authorYangqing Jia <jiayq84@gmail.com>
Tue, 17 Sep 2013 22:15:32 +0000 (15:15 -0700)
committerYangqing Jia <jiayq84@gmail.com>
Tue, 17 Sep 2013 22:15:32 +0000 (15:15 -0700)
src/caffeine/test/test_caffeine_main.cpp

index f76cc26cc72c882eb5305e2592c1781a5ae764b8..c4301054a2b3f64e822b5ddd1974f042f284db2b 100644 (file)
@@ -1,4 +1,5 @@
 #include <iostream>
+#include <cstdlib>
 
 #include <cuda_runtime.h>
 #include <glog/logging.h>
@@ -15,6 +16,12 @@ int main(int argc, char** argv) {
   int device;
   cudaGetDeviceCount(&device);
   cout << "Cuda number of devices: " << device << endl;
+  if (argc > 1) {
+    // Use the given device
+    device = atoi(argv[1]);
+    cudaSetDevice(device);
+    cout << "Setting to use device " << device << endl;
+  }
   cudaGetDevice(&device);
   cout << "Current device id: " << device << endl;
   cudaGetDeviceProperties(&CAFFEINE_TEST_CUDA_PROP, device);