I see. Thanks for the reply @ptrblck
I will read more about autocast tomorrow and will try tracing it.
Meanwhile, I also used the following script (using half) which gave me a fp16 graph
import torchvision
import torch
import tvm
from tvm import relay

is_fp16 = True

model = torchvision.models.resnet18(pretrained=True)
x = torch.rand(1, 3, 224, 224)
if is_fp16:
    model = model.cuda().half()
    x = x.cuda().half()
    scripted_model = torch.jit.trace(model, (x))
    scripted_model.save("rn50_fp16.pt")
else:
    scripted_model = torch.jit.trace(model, (x))
    scripted_model.save("rn50_fp32.pt")

input_name = "input0"
input_shape = (1, 3, 224, 224)
shape_list = [(input_name, input_shape)]
mod, params = relay.frontend.from_pytorch(scripted_model, shape_list)
mod = relay.transform.InferType()(mod)
print(mod)

My usecase is inference. The above code converted the whole model to fp16. But I am worried that it might be doing to much of fp16 and may need to go back to fp32 for some operators like sum or batch_norm.