add dimension check

Lightning-AI · Apr 23, 2024 · fb5071f · fb5071f
1 parent f7a6121
commit fb5071f
Showing 1 changed file with 7 additions and 4 deletions.
diff --git a/thunder/executors/nvfuserex_impl.py b/thunder/executors/nvfuserex_impl.py
@@ -2204,10 +2204,13 @@ def _matmul_check(
         return False
 
     enable_matmul: None | bool = get_compile_option("nv_enable_matmul", "Enable nvFuser matmul.")
-    if enable_matmul is None:
-        enable_matmul = False
-
-    return enable_matmul and is_supported_tensor(a) and is_supported_tensor(b)
+    if not enable_matmul:
+        return False
+    if not are_supported_tensors(a, b):
+        return False
+    if not (a.ndim == b.ndim and a.ndim == 2):
+        return False
+    return True
 
 
 def matmul(