Unverified Commit 72ca0f30 authored by ckirchhoff's avatar ckirchhoff Committed by GitHub
Browse files

[Fix] Fix the bug when input mask is not '0-1-tensor' in masked_conv (#2423)

parent 6e3827d0
...@@ -58,6 +58,8 @@ class MaskedConv2dFunction(Function): ...@@ -58,6 +58,8 @@ class MaskedConv2dFunction(Function):
if mask.size()[1:] != output.size()[2:]: if mask.size()[1:] != output.size()[2:]:
raise ValueError( raise ValueError(
'The mask is inconsistent with the shape of output_conv.') 'The mask is inconsistent with the shape of output_conv.')
mask = mask > 0
mask = mask.type(output.dtype)
output = output * mask output = output * mask
return output return output
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment