diff options
| author |  Anna R <annarev@google.com> | 2018-03-28 16:52:39 -0700 |
|---|---|---|
| committer |  TensorFlower Gardener <gardener@tensorflow.org> | 2018-03-28 16:55:15 -0700 |
| commit | 108178da2a20ea2d3899417ee932d46ba1a5c652 (patch) | |
| tree | 313bd8cec176f8c9ef67b25c6484a650d1f2092a /tensorflow/core/kernels/snapshot_op_gpu.cu.cc | |
| parent | 390e19ab990f5656e09d98624c92b3c80e52937d (diff) | |
Automated g4 rollback of changelist 190835392
PiperOrigin-RevId: 190858242
Diffstat (limited to 'tensorflow/core/kernels/snapshot_op_gpu.cu.cc')
| -rw-r--r-- | tensorflow/core/kernels/snapshot_op_gpu.cu.cc | 9 |
1 files changed, 6 insertions, 3 deletions
diff --git a/tensorflow/core/kernels/snapshot_op_gpu.cu.cc b/tensorflow/core/kernels/snapshot_op_gpu.cu.cc index e4e3bd5220..52070be838 100644 --- a/tensorflow/core/kernels/snapshot_op_gpu.cu.cc +++ b/tensorflow/core/kernels/snapshot_op_gpu.cu.cc @@ -24,10 +24,13 @@ limitations under the License. namespace tensorflow { typedef Eigen::GpuDevice GPUDevice; -// Definition of the GPU implementations declared in softsign_op.cc. -#define DEFINE_GPU_KERNELS(T) template struct functor::Snapshot<GPUDevice, T>; +#define REGISTER_KERNEL(TYPE) \ + REGISTER_KERNEL_BUILDER( \ + Name("Snapshot").Device(DEVICE_GPU).TypeConstraint<TYPE>("T"), \ + SnapshotOp<GPUDevice, TYPE>); -TF_CALL_POD_TYPES(DEFINE_GPU_KERNELS); +TF_CALL_POD_TYPES(REGISTER_KERNEL); +#undef REGISTER_KERNEL } // namespace tensorflow |