diff --git a/docs/how_to/tutorials/customize_opt.py b/docs/how_to/tutorials/customize_opt.py index 2e2747d61fc5..c175f43c6fcf 100644 --- a/docs/how_to/tutorials/customize_opt.py +++ b/docs/how_to/tutorials/customize_opt.py @@ -103,7 +103,16 @@ def forward(self, x): # Import cublas pattern -import tvm.relax.backend.cuda.cublas as _cublas +try: + import tvm.relax.backend.cuda.cublas as _cublas +except ImportError as e: + raise ImportError( + "This tutorial requires TVM built with CUDA support.\n" + "If you hit missing 'tvm_ffi', try: pip install apache-tvm-ffi\n" + "Otherwise build TVM with CUDA enabled:\n" + " https://tvm.apache.org/docs/install/from_source.html\n" + f"Original error: {e}" + ) from e # Define a new pass for CUBLAS dispatch @@ -174,7 +183,7 @@ def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IR # e.g. language model, DLight provides excellent performance, while for generic models, # it achieves a balance between performance and compilation time. -from tvm import dlight as dl +from tvm.s_tir import dlight as dl # Apply DLight rules with target: