I want to add preprocessing to a wrapper class around my forward model, and trace with TorchScript to be able to use in a C++ code. The normalization of input data, etc. works fine, but I want to be able to seamlessly handle inputs with different channel numbers (padding with 0’s if 1 channel). I tried doing this with if input.size()[0]==1, but get a warning that makes me think this might be ignored:

create_torchscript_model.py:37: TracerWarning: torch.tensor results are registered as constants in the trace. You can safely ignore this warning if you use this function to create tensors out of constant variables that would be the same every time you call this function. In any other case, this might cause the trace to be incorrect.

I try keeping everything in tensors:
        if torch.tensor(input.size())[0]==torch.tensor(1):
            input = F.pad(input,(0,0,0,0,0,0,1,0))

but get similar warnings. Is there a better way to do this, or should I just keep it in the C++ code?