The trainium toolchain is not as mature as GPU. Your model may fail to compile out of the box, and even if it does it may be slow and require you to dig into details for reasonable training/inference performance