int THCudaTensor_checkGPU(THCState *state, unsigned int nTensors, ...)
{
int kernelDev;
if (THCState_getDeviceMode(state) == THCStateDeviceModeManual) {
THCudaCheck(cudaGetDevice(&kernelDev));
} else {
kernelDev = THC_DEVICE_NONE;
}
va_list(args);
va_start(args, nTensors);
for (unsigned int i = 0; i < nTensors; i++) {
THCudaTensor* tensor = va_arg(args, THCudaTensor*);
if(tensor == NULL) {
continue;
}
int tensorDev = THCudaTensor_getDevice(state, tensor);
if (tensorDev != THC_DEVICE_NONE) {
if (kernelDev != tensorDev && kernelDev != THC_DEVICE_NONE) {
va_end(args);
return 0; // device mismatch
} else {
kernelDev = tensorDev;
}
}
}
va_end(args);
if (THCState_getDeviceMode(state) == THCStateDeviceModeAuto) {
if (kernelDev == THC_DEVICE_NONE) {
return 0; // cannot determine device
} else {
THCState_setDevice(state, kernelDev);
}
}
return 1;
}
static int cutorch_CudaTensor_getDevice(lua_State *L) {
THCudaTensor *tensor = luaT_checkudata(L, 1, "torch.CudaTensor");
lua_pushinteger(L, THCudaTensor_getDevice(cutorch_getstate(L), tensor) + 1);
return 1;
}