My guess is you may be seeing something some other of us have: Laptop Overheating and Losing Connection to M.2 Drive
The WD Black SN850 2 TB specifically can get too hot for the laptop and cause it to d/c. At least from what we’ve seen.
I will say that one of the first things I did was install the WD tool for the SSD and update the firmware. I seem to have less issues than others.