How to correctly use the "pad" function in TorchScript?

jit
1

The following code works fine in regular Pytorch

def pad_audio(self,tensor,n_max=50000):
    
    diff_pad = n_max - len(tensor)    
    
    tensor = F.pad(tensor,(int(diff_pad/2),diff_pad - int(diff_pad/2)),'constant',0)

But when using JIT script, I get the following error:


_pad(Tensor input, int[] pad, str mode="constant", float value=0.) -> (Tensor):
Expected a value of type 'List[int]' for argument 'pad' but instead found type 'Tuple[int, Tensor]'.
:
        diff_pad = n_max - len(tensor)    
        
        tensor = F.pad(tensor,(int(diff_pad/2),diff_pad - int(diff_pad/2)),'constant',0)
                 ~~~~~ <--- HERE

What is the correct way of using pad in TorchScript?

2

Make sure to use List[int] as described in the error message. Afterwards, you would run into another type mismatch, since the value is expected as a float and lastly you would have to annotate the inputs since n_max would be expected as a Tensor otherwise.
This should work:

@torch.jit.script
def pad_audio(tensor, n_max=50000):
    # type: (Tensor, int)  -> Tensor
    diff_pad = int(n_max - len(tensor))
    tensor = F.pad(tensor, [int(diff_pad/2), diff_pad - int(diff_pad/2)], 'constant', 0.)
    return tensor
1 Like
3

all parameters which are taken to F.pad() should have type integer, you can try this:

tensor = F.pad(tensor, [int(diff_pad/2), int(diff_pad) - int(diff_pad/2)], 'constant', 0.)