How to deal with prim::DictConstruct

masahi · June 14, 2022, 8:25pm

You can try something like

masahi/tvm-cutlass-eval/blob/master/bert/export.py#L26-L33


class TraceWrapper(torch.nn.Module):
    def __init__(self, model):
        super().__init__()
        self.model = model


    def forward(self, *inp):
        out = self.model(*inp)
        return out["logits"]