text_plain_output_1.png
import torch

model = MyModel()
model.eval()

x = torch.rand(1, 3, 336, 336, device='cuda')
model(x).shape

saved_model = torch.jit.script(model)
saved_model.save('saved_model.pt')

saved_model(x).shape