Use megatron-fsdp==0.1.0rc1 which has the >2D parameter fix.

cspades · cspades · commit a5023090e813 · 2025-09-02T15:11:21.000-07:00
Signed-off-by: Cory Ye &lt;cye@nvidia.com&gt;
diff --git a/recipes/vit/requirements.txt b/recipes/vit/requirements.txt
@@ -1,7 +1,7 @@
 torch
 torchvision
 transformer_engine
-megatron-fsdp==0.1.0rc0
+megatron-fsdp==0.1.0rc1
 hydra-core
 einops
 wandb
