Error in the space_to_dense function

I found a mistake in the space to dense function:

class space_to_dense(nn.Module):
    def __init__(self,stride):
        super(space_to_dense,self).__init__()
        self.stride = stride
    def forward(self, input):
        assert len(input.shape) == 4,"input tensor must be 4 dimenson"
        stride = self.stride
        B,C,W,H = input.shape
        assert (W %stride == 0 and H %stride == 0),"the W = {} or H = {} must be divided by {}".format(W,H,stride)
        ws = W // stride
        hs = H // stride
        x = input.view(B, C, hs, stride, ws, stride).transpose(3, 4).contiguous()
        x = x.view(B, C, hs*ws, stride * stride).transpose(2, 3).contiguous()
        x = x.view(B, C, stride * stride, hs, ws).transpose(1, 2).contiguous()
        x = x.view(B, stride * stride * C, hs, ws)
        return x

The line "B,C,W,H = input.shape" should be "B,C,H,W = input.shape"

If you load images having different height and width the concatenation between the lower and the higher order feature map complains about the different h and w of the two tensors.

lingtengqiu / Deeperlab-pytorch

Error in the space_to_dense function #6