fix(library): only compile CUDA extension on Linux

huggingface · Jan 6, 2025 · 24222c6 · 24222c6
1 parent 66dca87
commit 24222c6
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/optimum/quanto/library/extensions/__init__.py b/optimum/quanto/library/extensions/__init__.py
@@ -12,13 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import platform
+
 import torch
 
 from .cpp import *
 from .extension import *
 
 
-if torch.cuda.is_available():
+if torch.cuda.is_available() and platform.system() == "Linux":
     if torch.version.cuda:
         from .cuda import *
     elif torch.version.hip: