ReLU6 is not fused properly

[ilouzl]

Description of the bug:

When adding torch.nn.ReLU6 activation after Conv2D, the converter doesn't fuse it the Conv.
Only by tweaking Relu6 a little bit it is fused, but an additional Relu is appended to it.

Actual vs expected behavior:

The outcome of the conversion is a sequence of Conv2D -> Min -> Relu.
Expected output would be a single Conv2d fused with Relu6.

Any other information you'd like to share?

Sample code to reproduce:

import ai_edge_torch
import torch
import torchvision

model = torch.nn.Sequential(torch.nn.Conv2d(3, 32, 3), torch.nn.ReLU6())
sample_inputs = (torch.randn(1, 3, 64, 64),)
torch_output = model(*sample_inputs)
edge_model = ai_edge_torch.convert(model.eval(), sample_inputs)
edge_output = edge_model(*sample_inputs)
edge_model.export('relu6.tflite')


class MyRelu6(torch.nn.Module):
    def __init__(self, inplace=False):
        super(MyRelu6, self).__init__()
        self.inplace = inplace

    def forward(self, x):
        return torch.nn.functional.relu(x, inplace=self.inplace).clamp(0, 6)


model = torch.nn.Sequential(torch.nn.Conv2d(3, 32, 3), MyRelu6())
sample_inputs = (torch.randn(1, 3, 64, 64),)
torch_output = model(*sample_inputs)
edge_model = ai_edge_torch.convert(model.eval(), sample_inputs)
edge_output = edge_model(*sample_inputs)

edge_model.export('custom_relu6.tflite')

LEFT : relu6.tflite
RIGHT: custom_relu6.tflite: