Jongchan/attention-module

Bug in BAM: only support input height = width

HeMuling opened this issue · 0 comments

in BAM ChannelGate:

avg_pool = F.avg_pool2d( in_tensor, in_tensor.size(2), stride=in_tensor.size(2) )

this line assumes the input H=W

to fix this, on can change the code into:

avg_pool = F.avg_pool2d( in_tensor, (in_tensor.size(2), in_tensor.size(3)), stride=in_tensor.size(2) )