Hi,
I’m trying to create a quantized version of deeplab_v3_resnet50 model. I did the following changes in the Resnet and deeplab model class defintions

replaced addition with nn.quantized.FloatFunctional().add() in BottleNeck forward function inside resnet.py
added quant() and dequant() operations before and after the forward function definition in the deeplab_v3 model class.

But when creating an object of the above class, I get the following error : AttributeError: ‘Graph’ object has no attribute ‘_tracer_cls’.
The traceback is as follows:
~/Downloads/thinkAutonomous/modelOptimization/dl-optimization-repos/dl_scripts/dl-optimization/quantizedModel/deeplabv3Model.py in _deeplabv3_resnet(backbone, num_classes, aux)
    118     if aux:
    119         return_layers["layer3"] = "aux"
--> 120     backbone = create_feature_extractor(backbone, return_layers)
    121 
    122     aux_classifier = FCNHead(1024, num_classes) if aux else None

~/Downloads/thinkAutonomous/modelOptimization/dl-optimization-repos/dl_scripts/dl-optimization/quantizedModel/feature_extraction.py in create_feature_extractor(model, return_nodes, train_return_nodes, eval_return_nodes, tracer_kwargs, suppress_diff_warning)
    490 
    491     # Build the final graph module
--> 492     graph_module = DualGraphModule(model, graphs["train"], graphs["eval"], class_name=name)
    493 
    494     # Restore original training mode

~/Downloads/thinkAutonomous/modelOptimization/dl-optimization-repos/dl_scripts/dl-optimization/quantizedModel/feature_extraction.py in __init__(self, root, train_graph, eval_graph, class_name)
    261         # to re-create the Graph during deserialization.
    262         assert (
--> 263             self.eval_graph._tracer_cls == self.train_graph._tracer_cls
    264         ), "Train mode and eval mode should use the same tracer class"
    265         self._tracer_cls = None

AttributeError: 'Graph' object has no attribute '_tracer_cls'

From my understanding the graph = tracer.trace(model)  function in feature_extraction.py 3 returns a object of type nn.fx.Graph which doesn’t have a _tracer_cls function. I don’t understand how this error is related to the changes I made. By default, the deeplab_v3_resnet50 model works fine, but when I make these changes this error pops up.
It’d be helpful if someone could help me solve the issue.