add input index choice for mkldnn_concat

release/0.11.0
tensor-tang 7 years ago
parent 85b839f0f1
commit 9acfba82a3

@ -138,8 +138,11 @@ void MKLDNNLayer::backward(const UpdateCallback& callback) {
}
}
void MKLDNNLayer::reshapeInput(int& batchsize, int& height, int& width) {
const Argument& input = inputLayers_[0]->getOutput();
void MKLDNNLayer::reshapeInput(int& batchsize,
int& height,
int& width,
size_t inputIdx) {
const Argument& input = inputLayers_[inputIdx]->getOutput();
batchsize = input.getBatchSize();
int h = input.getFrameHeight();
int w = input.getFrameWidth();

@ -178,7 +178,10 @@ protected:
/**
* reshape the input image sizes and input batchsize
*/
void reshapeInput(int& batchsize, int& height, int& width);
void reshapeInput(int& batchsize,
int& height,
int& width,
size_t inputIdx = 0);
/**
* reshape output image sizes

Loading…
Cancel
Save