It can be made more efficient by always loading on CPU and only having FSDP2 handle the placement
It can be made more efficient by always loading on CPU and only having FSDP2 handle the placement