Relax the need for only CPU versions of PyTorch (#1505)

* Relax the need for only CPU versions of PyTorch

This allows installing corresponding PyTorch CUDA / ROCM versions and using torch-mlir.

* Remove obsolete comments

setup.py

@@ -155,11 +155,7 @@ setup(
     ],
     install_requires=[
         "numpy",
-        # To avoid issues with drift for each nightly build, we pin to the
-        # exact version we built against.
-        # TODO: This includes the +cpu specifier which is overly
-        # restrictive and a bit unfortunate.
-        f"torch=={torch.__version__}",
+        f"torch=={torch.__version__}".split("+", 1)[0],
     ],
     zip_safe=False,
 )