Corrected the tensor shape in the doc string. The incoming masks are stacked in dim=0 therefore the doc is wrong Signed-off-by: Jonas Heinle <jonasheinle@googlemail.com>
