Remove apex dependency after PyTorch 2.0?

Created by: suchenzang

This is to look into whether or not we still need to use apex for speedups if out-of-the-box PyTorch 2.0 may "just work". Will require benchmarking at a few different scales to confirm.