Thanks for your reply. Here’s the profiling code:
model = torch.load('path_to_model.pt')
model.eval()

for image in images:
    image = image.reshape(1, 3, 640, 640)
    start = time.time()
    res = model(torch.from_numpy(img).type(torch.cuda.FloatTensor) / 255)
    end = time.time()
    print(f'Run time: {end - start}')     # About 0.013s

    t0 = time.time()
    torch.cuda.synchronize()
    t1 = time.time()
    pred = res['out'].cpu().detach().numpy()[0][0]
    t2 = time.time()
    torch.cuda.synchronize()
    t3 = time.time()
    print(f'First sync time: {t1 - t0}')             # About 10s inside container, 0.08s outside
    print(f'CPU copy time: {t2 - t1}')               # About 6e-4s both in/out
    print(f'Second sync time: {t3 - t2}')          # About 3e-5s both in/out

Looks like I made a mistake and used torch.cuda.synchronize() in the wrong place previously. Now it seems CPU copy doesn’t take much time, but the first synchronization call takes very long inside the container.