Ensure custom GPU kernels are properly registered.
This change ensures that GOOGLE_CUDA=1 when compiling .cc files with --config=cuda enabled (also includes other important copts like -fno-exceptions.) Change: 129242223
This commit is contained in:
parent
ec73b4e819
commit
c5b3ea14c0
@ -29,11 +29,6 @@ limitations under the License.
|
||||
#include "tensorflow/core/util/padding.h"
|
||||
#include "tensorflow/core/util/tensor_format.h"
|
||||
|
||||
#if GOOGLE_CUDA
|
||||
#include "tensorflow/core/kernels/maxpooling_op_gpu.h"
|
||||
#include "tensorflow/core/kernels/pooling_ops_common_gpu.h"
|
||||
#endif // GOOGLE_CUDA
|
||||
|
||||
namespace tensorflow {
|
||||
|
||||
typedef Eigen::ThreadPoolDevice CPUDevice;
|
||||
|
@ -634,6 +634,7 @@ def tf_custom_op_library(name, srcs=[], gpu_srcs=[], deps=[]):
|
||||
srcs=srcs,
|
||||
deps=deps + if_cuda(cuda_deps),
|
||||
data=[name + "_check_deps"],
|
||||
copts=tf_copts(),
|
||||
linkshared=1,
|
||||
linkopts = select({
|
||||
"//conditions:default": [
|
||||
|
@ -1 +1,2 @@
|
||||
*tensorflow*
|
||||
*perftools*gputools*
|
||||
|
@ -1,6 +1,7 @@
|
||||
tensorflow {
|
||||
global:
|
||||
*tensorflow*;
|
||||
*perftools*gputools*;
|
||||
local:
|
||||
*;
|
||||
};
|
||||
|
Loading…
Reference in New Issue
Block a user