Hi @cfanyyx, thanks for raising the issue. The problem on this is because nn.Upsample is a module that contains data dependent control flow (which means the traced model totally depends on the input you provide during tracing and will not generalize to the new input). To make nn.Upsample working in JIT, you will need to use TorchScript scripting instead of tracing. You can refer more to the doc here 47
Let me know if that’s clear or not.