Add compile guard to Dropout's FIsCUDAGraphsCompatible def

DickJC123 · Feb 18, 2022 · c44cfc6 · c44cfc6
1 parent eaa7fc7
commit c44cfc6
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/src/operator/nn/dropout.cu b/src/operator/nn/dropout.cu
@@ -34,14 +34,18 @@ NNVM_REGISTER_OP(Dropout)
                                          // Dropout is a passthrough during inference for all impls
                                          if (!is_train)
                                            return true;
+#if MXNET_USE_CUDNN_DROPOUT
                                          // cuDNN impl is compatible during training as well
                                          const DropoutParam& param =
                                              nnvm::get<DropoutParam>(attrs.parsed);
                                          real_t pkeep = 1.0f - param.p;
                                          bool cudnn_off =
                                              param.cudnn_off && param.cudnn_off.value();
                                          bool cudnn_available = pkeep > 0 && !cudnn_off;
-                                         return MXNET_USE_CUDNN_DROPOUT && cudnn_available;
+                                         return cudnn_available;
+#else
+                                         return false;
+#endif  // MXNET_USE_CUDNN_DROPOUT
                                        })
     .set_attr<FStatefulCompute>("FStatefulCompute<gpu>", DropoutCompute<gpu>);