🐛 [Bug] torch.export fails to export embedding_bag_aten layer in a standalone test

[peri044]

Bug Description

Error message:

torch._dynamo.exc.UserError: Cannot export model which references tensors that are neither buffers/parameters/constants nor are direct inputs.  For each tensor, if you'd like this tensor to be an explicit input, add it as a dummy argument to the top-level model definition you are exporting; if you would like its value to be embedded as an exported constant, wrap its access in a function marked with @assume_constant_result.
E           
E           L['offsets'], a closed over free variable
E           
E           To execute this test, run the following from the base repo dir:
E                python test_embedding_bag_aten.py -k test_embedding_bag_with_traversable_offsets_7

Currently we are using torch.fx.symbolic_trace(mod) + ShapeProp to make it pass (this was how it was originally written for). ShapeProp is needed to populate metadata which is being used in capability validation otherwise the metadata is empty which results in Conversion of function torch._ops.aten.aten::_embedding_bag not currently supported! issue

To Reproduce

Steps to reproduce the behavior:

Expected behavior

Environment

Build information about Torch-TensorRT can be found by turning on debug messages

• Torch-TensorRT Version (e.g. 1.0.0):
• PyTorch Version (e.g. 1.0):
• CPU Architecture:
• OS (e.g., Linux):
• How you installed PyTorch (conda, pip, libtorch, source):
• Build command you used (if compiling from source):
• Are you using local sources or building from archives:
• Python version:
• CUDA version:
• GPU models and configuration:
• Any other relevant information:

Additional context

[zewenli98]

@peri044 Can you provide a reproducer? I tried something like below that works fine on my end:

class TestEmbeddingBag(torch.nn.Module):
    def forward(self, weight, indices, offsets):
        return torch.ops.aten._embedding_bag.default(
            weight,
            indices,
            offsets,
            scale_grad_by_freq,
            mode,
            sparse,
            per_sample_weights,
            include_last_offset,
            padding_idx,
        )[0]

model = TestEmbeddingBag()
exported_program = torch.export.export(model, (weight, indices, offsets))
print(exported_program)

[zewenli98]

Currently we are using torch.fx.symbolic_trace(mod) + ShapeProp to make it pass (this was how it was originally written for).

What's the new way to populate metadata? I'm not familiar with this part