Fix PyTorch version compatibility check in _load_library

gchalump · spcyppt · commit bc3b58a08faf · 2025-10-21T11:22:41.000-07:00
The version compatibility check in `_load_library` incorrectly uses exact string equality when comparing PyTorch versions, causing false incompatibility warnings.

PyTorch versions typically include build metadata like `"2.8.0+cu118"` or `"2.8.1+cpu"`, while the compatibility table contains base versions like `"2.8"`.
diff --git a/fbgemm_gpu/fbgemm_gpu/__init__.py b/fbgemm_gpu/fbgemm_gpu/__init__.py
@@ -81,7 +81,7 @@ def _load_library(filename: str, version: str, no_throw: bool = False) -> None:
             """
         )
 
-    elif str(torch.__version__) != _fbgemm_torch_compat_table[keys[0]]:
+    elif not str(torch.__version__).startswith(_fbgemm_torch_compat_table[keys[0]]):
         logging.warning(
             f"""
             \033[31m

Original file line number	Diff line number	Diff line change
`@@ -81,7 +81,7 @@ def _load_library(filename: str, version: str, no_throw: bool = False) -> None:`
`81`	`81`	`"""`
`82`	`82`	`)`
`83`	`83`
`84`		`- elif str(torch.__version__) != _fbgemm_torch_compat_table[keys[0]]:`
	`84`	`+ elif not str(torch.__version__).startswith(_fbgemm_torch_compat_table[keys[0]]):`
`85`	`85`	`logging.warning(`
`86`	`86`	`f"""`
`87`	`87`	`\033[31m`