Pipeline parallelism examples with Pippy fails

[goelayu]

System Info

- `Accelerate` version: 0.35.0.dev0
- Platform: Linux-5.15.0-121-generic-x86_64-with-glibc2.35
- `accelerate` bash location: redacted
- Python version: 3.10.14
- Numpy version: 1.23.5
- PyTorch version (GPU?): 2.4.1+cu121 (True)
- PyTorch XPU available: False
- PyTorch NPU available: False
- PyTorch MLU available: False
- PyTorch MUSA available: False
- System RAM: 1007.59 GB
- GPU type: NVIDIA H100 PCIe

Information

• [X] The official example scripts
• [ ] My own modified scripts

Tasks

• [ ] One of the scripts in the examples/ folder of Accelerate or an officially supported no_trainer script in the examples folder of the transformers repo (such as run_no_trainer_glue.py)
• [ ] My own task or dataset (give details below)

Reproduction

1. Run the llama.py example in distributed inference folder.
2. Getting the following error

   torch._dynamo.exc.UserError: Dynamic control flow is not supported at the moment. Please use functorch.experimental.control_flow.cond to explicitly capture the control flow. 
   For more information about this error, see: https://pytorch.org/docs/main/generated/exportdb/index.html#cond-operands

Expected behavior

I have tried using different accelerate launch flags such as --dynamo_use_dynamic, however I am not sure how to fix the above error.

[muellerzr]

@goelayu can you try upgrading your python version? IIRC that can play a role. (3.12 ideally)