for i in range(num_steps):
    # This line might take 5-10ms per image, done sequentially
    batch = load_and_preprocess_batch(file_paths[i])
    
    # This line might take 0.5ms thanks to JIT
    loss = train_step(model, optimizer, batch)

        __
        __