[rocm] Fix current torch.linalg.lstsq limitation in InverseMelScale module #3939
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This PR fixes the following error:
RuntimeError: torch.linalg.lstsq: only overdetermined systems (input.size(-2) >= input.size(-1)) are allowed on CUDA. Please rebuild with cuSOLVER
Checking the pytorch code https://github.com/ROCm/pytorch/blob/main/aten/src/ATen/native/cuda/linalg/BatchLinearAlgebra.cpp, it calls magma. According to magma documentation, least squares method - magma_sgels_gpu supports only overdetermined systems https://icl.utk.edu/projectsfiles/magma/doxygen/group__magma__gels.html
The proposed temporary solution uses https://docs.pytorch.org/docs/stable/generated/torch.linalg.pinv.html as a replacement for linear least squares. This change can be reverted once torch.linalg.lstsq uses hipSolver for underconstrained systems.
Tested on rocm with
pytest test/torchaudio_unittest/transforms/transforms_cuda_test.py -k test_inverse_melscale