experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Merge pull request #33021 from fo40225:fix-msvc163-cuda101

Browse Source 
PiperOrigin-RevId: 276092173
Change-Id: Id8be850452bb3c4f93a212eb36e48458d24052c0
This commit is contained in:
TensorFlower Gardener 2019-10-22 11:07:01 -07:00
commit e9597468d8
1 changed files with 4 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
tensorflow/core/kernels/scatter_nd_op_gpu.cu.cc
View File

@ -64,7 +64,7 @@ struct LeftUpdate<std::complex<T>, scatter_nd_op::UpdateOp::ADD> {
std::complex<T>* out, const std::complex<T>& val) { std::complex<T>* out, const std::complex<T>& val) {
T* ptr = reinterpret_cast<T*>(out); T* ptr = reinterpret_cast<T*>(out);
GpuAtomicAdd(ptr, val.real()); GpuAtomicAdd(ptr, val.real());
GpuAtomicAdd(ptr, val.imag()); GpuAtomicAdd(ptr + 1, val.imag());
} }
}; };
@ -72,7 +72,9 @@ template <typename T>
struct LeftUpdate<std::complex<T>, scatter_nd_op::UpdateOp::SUB> { struct LeftUpdate<std::complex<T>, scatter_nd_op::UpdateOp::SUB> {
EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC void operator()( EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC void operator()(
std::complex<T>* out, const std::complex<T>& val) { std::complex<T>* out, const std::complex<T>& val) {
LeftUpdate<std::complex<T>, scatter_nd_op::UpdateOp::ADD>()(out, -val); T* ptr = reinterpret_cast<T*>(out);
GpuAtomicSub(ptr, val.real());
GpuAtomicSub(ptr + 1, val.imag());
} }
}; };