I have finetuned DeepLabv3 model for a custom dataset.
I converted the trained model to a jit model using torch.jit.script() and the model worked fine and gave the desired output during inference. But when I try to convert the same model again to jit, the results are undesirable during inference and looks completely poles apart.
I am new to jit, so trying to figure out things.
The model was trained on 2 GPUs, then convert to a single model GPU, then tried converting it to jit.
Both the outputs are before converting to a jit model. So, I think converting from multi gpu to single gpu is causing the issue. Is there a way to convert a multigpu model to single gpu other than the above mentioned way?
Using the above mentioned way to convert nn.Dataparallel to single gpu model works consistently compared to the key replacement in the state_dict. Now I am able to get desired solution and indeed the issue is with the state_dict as foreseen by @ptrblck