Use a CUDAGuard when running Torch models

Author: VivekPanyam

source/neuropod/backends/torchscript/torch_backend.cc

@@ -7,6 +7,7 @@
 #include "neuropod/backends/torchscript/type_utils.hh"
 #include "neuropod/internal/tensor_types.hh"
 
+#include <c10/cuda/CUDAGuard.h>
 #include <caffe2/core/macros.h>
 
 #include <iostream>
@@ -291,6 +292,14 @@ std::unique_ptr<NeuropodValueMap> TorchNeuropodBackend::infer_internal(const Neu
 {
     torch::NoGradGuard guard;
 
+    // Make sure we're running on the correct device
+    std::unique_ptr<at::cuda::CUDAGuard> device_guard;
+    const auto                           model_device = get_torch_device(DeviceType::GPU);
+    if (model_device.is_cuda())
+    {
+        device_guard = stdx::make_unique<at::cuda::CUDAGuard>(model_device);
+    }
+
     // Get inference schema
     const auto &method    = model_->get_method("forward");
     const auto &schema    = SCHEMA(method);