How To Choose Correct PyTorch Install in Python setup.py File?

1

Hello, I am working on a Python package that uses Pytorch as a dependency. In the setup.py file for the package, torch is listed as a dependency. Currently, the code looks something like this:

from distutils.core import setup

setup(
    name='cool package',
    description='A really cool package',
    author='me',
    url='https://my_cool_website.com',
    packages=['my_cool_package'],
    python_requires='>=3.7.10',
    install_requires=[
            'torch'
        ]
)

Currently, this will install the CPU compatible version of Pytorch. In the installation instructions on the Pytorch homepage, however, the following command is listed as the installation instructions to install Pytorch with GPU support for Cuda 10.2:

pip3 install torch==1.10.1+cu102 torchvision==0.11.2+cu102 torchaudio===0.10.1+cu102 -f https://download.pytorch.org/whl/cu102/torch_stable.html

Is there a way to make sure that I install this version of Pytorch in the setup.py file if the user has GPU support on their machine?