Tracing models with multi batch input

[zobertke]

Hi folks 👋 I have a model with an input of a list of 3 disconnected graphs represented as torch_geometric.data.Batch-es like:

example_batch = (
            Batch.from_data_list([graph1]),
            Batch.from_data_list([graph2]),
            Batch.from_data_list([graph3]),
)
traced_model = torch.jit.trace(model, example_batch)

I would like to trace this model with jit.trace and I can't really get around to feed it with an example data like that. Is there any support for this in the current version?
Thanks for your help

[rusty1s]

jit.trace does only support Tensors as input, so I suggest to change the forward arguments from your model from

def forward(self, data: List[Data])

to

def forward(self, x: List[Tensor], edge_index: List[Tensor])

[zobertke]

So I am assuming I need to create an adapter module, which does the transform front and back before/after passing it to the model. Thank you!

[rusty1s]

What do you mean by that? Do you desperately need access to Data inside your forward function?