fix

flashinfer-ai · Jan 31, 2024 · 4629a1c · 4629a1c
1 parent d5f397d
commit 4629a1c
Showing 1 changed file with 12 additions and 1 deletion.
diff --git a/python/flashinfer/utils.py b/python/flashinfer/utils.py
@@ -13,7 +13,18 @@
 See the License for the specific language governing permissions and
 limitations under the License.
 """
-import torch
+try:
+    import torch
+except ImportError as e:
+    import os
+    import logging
+    if os.environ.get("BUILD_DOC", "0") == "1":
+        _kernels = None
+        logging.warning(
+            "torch is not loaded in documentation build mode."
+        )
+    else:
+        raise e
 
 
 class RotaryMode: