Makes sense. So, I followed up on autocast. This is what my code looks like
model = torchvision.models.resnet18()

x = torch.rand(1, 3, 224, 224).cuda()
model = model.cuda()
with torch.cuda.amp.autocast(enabled=True):
    scripted_model = torch.jit.trace(model, (x))
    scripted_model.save("rn50_amp.pt")

It failed with the following error - RuntimeError: Cannot insert a Tensor that requires grad as a constant. Consider making it a parameter or input, or detaching the gradient
I then added no_grad
model = torchvision.models.resnet18()

x = torch.rand(1, 3, 224, 224).cuda()
model = model.cuda()
with torch.cuda.amp.autocast(enabled=True):
    with torch.no_grad():
        scripted_model = torch.jit.trace(model, (x))
        scripted_model.save("rn50_amp.pt")

But this also failed with some different error
torch.jit._trace.TracingCheckError: Tracing failed sanity checks!
ERROR: Graphs differed across invocations!

I am not sure if my usage is correct.