I’m aware that this isn’t a straight forward “buy this” answer, but it gives you insight what to watch out for when choosing an SSD for your specific use case.
A simpler “answer some questions, get a list” method is also on their main site (which has also been posted in the thread eray linked).
I’m not sure if you can still find any (Any that aren’t sold used with the code already taken, heh!), but if you get the original Optane 900p it came with a free Star Citizen Sabre Raven ship code.
I believe that the Optane NVMe drives are the absolute best SSD available. They’re really pricey though.
I went with a 960 GB 905p for my software development Ryzen 3900X Linux build. It’s amazingly good. I am constantly low on space though, and wish I’d gone for the 1.5 TB version. But, I’m going to be adding a PCIe 4 NVMe drive when it arrives from Amazon. They had a good deal.
Optane is a pretty great choice for ML if that’s what your system will be doing most of the time, then I see no reason why not, if it won’t be the primary use case, then any of the flagship SSDs even with Gen 3 speeds will do great for most things