anirudh2290 commented on a change in pull request #10014: Fix crash with 
mx.nd.ones
URL: https://github.com/apache/incubator-mxnet/pull/10014#discussion_r172751325
 
 

 ##########
 File path: src/imperative/imperative_utils.h
 ##########
 @@ -70,12 +70,14 @@ inline Context GetContext(const nnvm::NodeAttrs& attrs,
   if (ctx.dev_mask() != ctx.dev_type) {
     ctx = Context::Create(ctx.dev_mask(), ctx.dev_id);
   }
-#if !MXNET_USE_CUDA
   if (ctx.dev_mask() == gpu::kDevMask) {
+#if !MXNET_USE_CUDA
     LOG(INFO) << "GPU support is disabled. Compile MXNet with "
               << "USE_CUDA=1 to enable GPU support.";
-  }
+#else
+    mshadow::SetDevice<gpu>(ctx.dev_id);
 
 Review comment:
   The call is intended to check if the device_id is available. The overhead 
should be pretty low as it is only setting the dev id corresponding to the 
thread and there is no synchronization taking place. The other option is to get 
the cuda device count and check if it the device falls in the range but that 
would have a similar overhead. 

----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on GitHub and use the
URL above to go to the specific comment.
 
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org


With regards,
Apache Git Services

Reply via email to