Number of output channels of first arg do not match number input channels of second arg

Raise code

''' ranges of the first tensors' output channel are equal to the
    ranges of the second tensors' input channel
    '''
    if type(module1) not in _supported_types or type(module2) not in _supported_types:
        raise ValueError("module type not supported:", type(module1), " ", type(module2))

    if module1.weight.size(output_axis) != module2.weight.size(input_axis):
        raise TypeError("Number of output channels of first arg do not match \
        number input channels of second arg")

    weight1 = module1.weight
    weight2 = module2.weight
    bias = module1.bias

    weight1_range = channel_range(weight1, output_axis)
    weig