This is the code I used to support ConvBnReLU6 and ConvReLU6
def fuse_model(model):
    for m in model.modules():
        if type(m) == DepthwiseSeparableConv:
            torch.quantization.fuse_modules(m, ['conv_dw', 'bn1', 'act1'], inplace=True,
                                            fuser_func=fuse_known_modules_mod)
            # torch.quantization.fuse_modules(m, ['conv_pw', 'bn2', 'act2'], inplace=True,
            torch.quantization.fuse_modules(m, ['conv_pw', 'bn2'], inplace=True,
                                            fuser_func=fuse_known_modules_mod)
        elif type(m) == InvertedResidual:
            torch.quantization.fuse_modules(m, ['conv_pw', 'bn1', 'act1'], inplace=True,
                                            fuser_func=fuse_known_modules_mod)
            torch.quantization.fuse_modules(m, ['conv_dw', 'bn2', 'act2'], inplace=True,
                                            fuser_func=fuse_known_modules_mod)
            torch.quantization.fuse_modules(m, ['conv_pwl', 'bn3'], inplace=True, fuser_func=fuse_known_modules_mod)
    torch.quantization.fuse_modules(model, ['conv_head', 'bn2', 'act2'], inplace=True,
                                    fuser_func=fuse_known_modules_mod)
    torch.quantization.fuse_modules(model, ['conv_stem', 'bn1', 'act1'], inplace=True,
                                    fuser_func=fuse_known_modules_mod)


def fuse_known_modules_mod(mod_list):
    r"""Returns a list of modules that fuses the operations specified
     in the input module list.
    Fuses only the following sequence of modules:
    conv, bn
    conv, bn, relu
    conv, relu
    linear, relu
    For these sequences, the first element in the output module list performs
    the fused operation. The rest of the elements are set to nn.Identity()
    """

    OP_LIST_TO_FUSER_METHOD = {
        (torch.nn.Conv2d, torch.nn.BatchNorm2d): fuse_conv_bn,
        (torch.nn.Conv2d, torch.nn.BatchNorm2d, torch.nn.ReLU): fuse_conv_bn_relu,
        (torch.nn.Conv2d, torch.nn.BatchNorm2d, torch.nn.ReLU6): fuse_conv_bn_relu6,
        (Conv2dSame, torch.nn.BatchNorm2d, torch.nn.ReLU6): fuse_conv_bn_relu6,
        (torch.nn.Conv2d, torch.nn.ReLU): torch.nn.intrinsic.ConvReLU2d,
        (torch.nn.Conv2d, torch.nn.ReLU6): ConvReLU6,
        (torch.nn.Linear, torch.nn.ReLU): torch.nn.intrinsic.LinearReLU
    }

    types = tuple(type(m) for m in mod_list)
    fuser_method = OP_LIST_TO_FUSER_METHOD.get(types, None)
    if fuser_method is None:
        raise NotImplementedError("Cannot fuse modules: {}".format(types))
    new_mod = [None] * len(mod_list)
    new_mod[0] = fuser_method(*mod_list)

    for i in range(1, len(mod_list)):
        new_mod[i] = torch.nn.Identity()
        new_mod[i].training = mod_list[0].training

    return new_mod


class ConvReLU6(nn.Sequential):
    def __init__(self, conv, relu6):
        super(ConvReLU6, self).__init__(conv, relu6)


class ConvBnReLU6(torch.nn.Sequential):
    def __init__(self, conv, bn, relu6):
        super(ConvBnReLU6, self).__init__(conv, bn, relu6)


def fuse_conv_bn_relu6(conv, bn, relu6):
    assert (conv.training == bn.training == relu6.training), \
        "Conv and BN both must be in the same mode (train or eval)."

    if conv.training:
        return ConvBnReLU6(conv, bn, relu6)
    else:
        return ConvReLU6(
            torch.nn.utils.fusion.fuse_conv_bn_eval(conv, bn), relu6)