TorchScript hardcoded tensor shapes

jit
1

Hello !

Let’s say I have a warping module that inherits from nn.Module.
In the constructor of this module I create a meshgrid of a particular shape as a registered buffer.
Then in the forward method of this module, which takes a batch of images and a batch of homographies as inputs, I compute a new grid of coordinates using input homographies and the base meshgrid defined in the constructor. Then I use grid_sample to warp input images.

When I use torch.jit.trace() on this module, it seems only usable with the size given to torch.jit.trace(), “input_sample”. So the .pt export seems not dynamic, the size of the meshgrid in the constructor seems hardcoded in the TorchScript.

How can I do to have a “dynamic” TorchScript that would allow me to warp any shapes ?

Thanks for your help !

2

you want to check up torch.jit.script() instead - also keep in mind torchscript is in maintenance mode so compiler efforts are mostly under torch.compile()now which has a dynamic parameter you can use

3

Ok thank you I will check this out !