@@ -298,11 +298,13 @@ __global__ void KeAvgPoolBackward(const int nthreads,
298298
299299 for (int ph = phstart; ph < phend; ++ph) {
300300 int hstart = ph * strideH - padH;
301- int hend = min (hstart + sizeY, height + padH);
301+ int hend = min (hstart + sizeY, height);
302+ hstart = max (hstart, 0 );
302303 for (int pw = pwstart; pw < pwend; ++pw) {
303304 // figure out the pooling size
304305 int wstart = pw * strideW - padW;
305- int wend = min (wstart + sizeX, width + padW);
306+ int wend = min (wstart + sizeX, width);
307+ wstart = max (wstart, 0 );
306308 int poolsize = (hend - hstart) * (wend - wstart);
307309 gradient += outGrad[ph * pooledW + pw] / poolsize;
308310 }
@@ -708,14 +710,17 @@ __global__ void KeAvgPool3DBackward(const int nthreads,
708710
709711 for (int pd = pdstart; pd < pdend; ++pd) {
710712 int dstart = pd * strideD - padD;
711- int dend = min (dstart + sizeZ, depth + padD);
713+ int dend = min (dstart + sizeZ, depth);
714+ dstart = max (dstart, 0 );
712715 for (int ph = phstart; ph < phend; ++ph) {
713716 int hstart = ph * strideH - padH;
714- int hend = min (hstart + sizeY, height + padH);
717+ int hend = min (hstart + sizeY, height);
718+ hstart = max (hstart, 0 );
715719 for (int pw = pwstart; pw < pwend; ++pw) {
716720 // figure out the pooling size
717721 int wstart = pw * strideW - padW;
718- int wend = min (wstart + sizeX, width + padW);
722+ int wend = min (wstart + sizeX, width);
723+ wstart = max (wstart, 0 );
719724 int poolsize = (dend - dstart) * (hend - hstart) * (wend - wstart);
720725 gradient += outGrad[(pd * pooledH + ph) * pooledW + pw] / poolsize;
721726 }
0 commit comments