Skip to content

Commit f540233

Browse files
WindQAQseanpmorgan
authored andcommitted
Build correlation cost GPU kernel (#459)
* build gpu kernel * remove unused header * revert some changes * lint
1 parent cc89403 commit f540233

File tree

1 file changed

+1
-2
lines changed
  • tensorflow_addons/custom_ops/layers

1 file changed

+1
-2
lines changed

tensorflow_addons/custom_ops/layers/BUILD

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,14 +10,13 @@ cc_binary(
1010
srcs = [
1111
"cc/kernels/correlation_cost_op.cc",
1212
"cc/kernels/correlation_cost_op.h",
13-
"cc/kernels/correlation_cost_op_gpu.cu.cc",
1413
"cc/ops/correlation_cost_op.cc",
1514
],
1615
copts = [
1716
"-pthread",
1817
"-std=c++11",
1918
D_GLIBCXX_USE_CXX11_ABI,
20-
],
19+
] + if_cuda(["-DGOOGLE_CUDA=1"]),
2120
linkshared = 1,
2221
deps = [
2322
"@local_config_tf//:libtensorflow_framework",

0 commit comments

Comments
 (0)