mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-30 20:18:08 +00:00
Fix 3828 (#3837)
This commit is contained in:
parent
f0c9fbc051
commit
b45ce92542
1 changed files with 28 additions and 28 deletions
|
|
@ -166,40 +166,40 @@ Status Conv<T>::ComputeInternal(OpKernelContext* context) const {
|
|||
s_.algo = perf.algo;
|
||||
s_.workspace_bytes = perf.memory;
|
||||
}
|
||||
}
|
||||
|
||||
if (!y_data) {
|
||||
Tensor* Y = context->Output(0, TensorShape(s_.y_dims));
|
||||
// special case when there is a dim value of 0 in the shape.
|
||||
if (Y->Shape().Size() == 0)
|
||||
return Status::OK();
|
||||
if (!y_data) {
|
||||
Tensor* Y = context->Output(0, TensorShape(s_.y_dims));
|
||||
// special case when there is a dim value of 0 in the shape.
|
||||
if (Y->Shape().Size() == 0)
|
||||
return Status::OK();
|
||||
|
||||
y_data = reinterpret_cast<CudaT*>(Y->template MutableData<T>());
|
||||
}
|
||||
y_data = reinterpret_cast<CudaT*>(Y->template MutableData<T>());
|
||||
}
|
||||
|
||||
const auto alpha = Consts<CudaT>::One;
|
||||
const auto beta = Consts<CudaT>::Zero;
|
||||
const auto alpha = Consts<CudaT>::One;
|
||||
const auto beta = Consts<CudaT>::Zero;
|
||||
|
||||
IAllocatorUniquePtr<void> workspace = GetScratchBuffer<void>(s_.workspace_bytes);
|
||||
IAllocatorUniquePtr<void> workspace = GetScratchBuffer<void>(s_.workspace_bytes);
|
||||
|
||||
CUDNN_RETURN_IF_ERROR(cudnnConvolutionForward(CudnnHandle(),
|
||||
&alpha,
|
||||
s_.x_tensor,
|
||||
x_data,
|
||||
s_.filter_desc,
|
||||
w_data,
|
||||
s_.conv_desc,
|
||||
s_.algo,
|
||||
workspace.get(),
|
||||
s_.workspace_bytes,
|
||||
&beta,
|
||||
s_.y_tensor,
|
||||
y_data));
|
||||
CUDNN_RETURN_IF_ERROR(cudnnConvolutionForward(CudnnHandle(),
|
||||
&alpha,
|
||||
s_.x_tensor,
|
||||
x_data,
|
||||
s_.filter_desc,
|
||||
w_data,
|
||||
s_.conv_desc,
|
||||
s_.algo,
|
||||
workspace.get(),
|
||||
s_.workspace_bytes,
|
||||
&beta,
|
||||
s_.y_tensor,
|
||||
y_data));
|
||||
|
||||
if (has_bias) {
|
||||
const Tensor* B = context->Input<Tensor>(2);
|
||||
auto b_data = reinterpret_cast<const CudaT*>(B->template Data<T>());
|
||||
CUDNN_RETURN_IF_ERROR(cudnnAddTensor(CudnnHandle(), &alpha, s_.b_tensor, b_data, &alpha, s_.y_tensor, y_data));
|
||||
if (has_bias) {
|
||||
const Tensor* B = context->Input<Tensor>(2);
|
||||
auto b_data = reinterpret_cast<const CudaT*>(B->template Data<T>());
|
||||
CUDNN_RETURN_IF_ERROR(cudnnAddTensor(CudnnHandle(), &alpha, s_.b_tensor, b_data, &alpha, s_.y_tensor, y_data));
|
||||
}
|
||||
}
|
||||
|
||||
return Status::OK();
|
||||
|
|
|
|||
Loading…
Reference in a new issue