cudnn5 branch for tx2

This commit is contained in:
Francesco Gatti
2017-07-14 11:11:51 +02:00
parent 33513468e6
commit e4df86a07c
9 changed files with 49 additions and 276 deletions
+3 -3
View File
@@ -29,7 +29,7 @@ Conv2d::Conv2d( Network *net, dataDim_t in_dim, int out_ch,
net->tensorFormat, net->dataType, n, c, h, w) );
checkCUDNN( cudnnSetFilter4dDescriptor(filterDesc,
net->dataType, out_ch, input_dim.c,
net->dataType, net->tensorFormat, out_ch, input_dim.c,
kernelH, kernelW) );
checkCUDNN( cudnnSetConvolution2dDescriptor(convDesc,
@@ -103,7 +103,7 @@ value_type* Conv2d::infer(dataDim_t &dim, value_type* srcData) {
// bias
alpha = value_type(1);
beta = value_type(1);
checkCUDNN( cudnnAddTensor(net->cudnnHandle, CUDNN_ADD_SAME_C,
checkCUDNN( cudnnAddTensor(net->cudnnHandle,
&alpha, biasTensorDesc, bias_d,
&beta, dstTensorDesc, dstData) );
@@ -113,4 +113,4 @@ value_type* Conv2d::infer(dataDim_t &dim, value_type* srcData) {
return dstData;
}
}
}