@albanD : I am passing “y” to the torch.jit.trace, y has size = torch.Size([1, 3, 800, 800])
and if I pass the “image” variable to jit.trace, then image has size = torch.Size([3, 800, 800])
Also, please keep noted that error is same for both cases, if I pass “y” or “image” to the jit.trace
P.S: if I run the code without tracing, it works fine and the input image has size torch.Size([3, 800, 800])