[GML] Update setup.py with -gml version

Signed-off-by: James Bartlett <[email protected]>

CMakeLists.txt

@@ -34,6 +34,8 @@ include(CMakeDependentOption)
 option(TORCH_MLIR_ENABLE_WERROR_FLAG "Enable `-Werror` flag on supported directories, treat error as warning" OFF)
 option(TORCH_MLIR_USE_INSTALLED_PYTORCH "If depending on PyTorch use it as installed in the current Python environment" ON)
 
+set(TORCH_CUDA_ARCH_LIST "8.0 8.6 8.9 9.0")
+
 option(TORCH_MLIR_ENABLE_REFBACKEND "Enable reference backend" ON)
 if(TORCH_MLIR_ENABLE_REFBACKEND)
   add_definitions(-DTORCH_MLIR_ENABLE_REFBACKEND)

setup.py

@@ -223,16 +223,17 @@ def build_extension(self, ext):
 EXT_MODULES = [
     CMakeExtension("torch_mlir._mlir_libs._torchMlir"),
 ]
-NAME = "torch-mlir-core"
+NAME = "torch-mlir-core-gml"
 
 # If building PyTorch extensions, customize.
 if not TORCH_MLIR_ENABLE_ONLY_MLIR_PYTHON_BINDINGS:
     import torch
 
-    NAME = "torch-mlir"
+    NAME = "torch-mlir-gml"
     INSTALL_REQUIRES.extend(
         [
-            f"torch=={torch.__version__}".split("+", 1)[0],
+            # Don't explicitly depend on a torch version.
+            "torch",
         ]
     )
     EXT_MODULES.extend(