@richard Autograd is one issue, I can manually compute the gradients. But another issue is the communication.
Every time I need to do a svd or symeig, I have to copy the matrix from GPU to CPU and convert it to Cupy array and send to GPU back. Then after the actual operation, I copy it back to CPU and convert it to pytorch tensor and send back to GPU again. Using Cupy introduces this two round-trip communication, which is a little problematic.