Fix docs

lightx2v/common/ops/attn/attn_weight.py

@@ -23,7 +23,7 @@ except ImportError:
     logger.info("flash_attn_varlen_func_v3 not found, please install flash_attn3 first")
     flash_attn_varlen_func_v3 = None
 
-if torch.cuda.get_device_capability(0)[0] <= 8 and torch.cuda.get_device_capability(0)[1] <= 9:
+if torch.cuda.get_device_capability(0) == (8, 9):
     try:
         from sageattention import sageattn_qk_int8_pv_fp16_triton as sageattn
     except ImportError:

lightx2v/models/networks/wan/infer/transformer_infer.py

@@ -29,6 +29,8 @@ class WanTransformerInfer(BaseTransformerInfer):
         self.mask_map = None
 
         if self.config["cpu_offload"]:
+            if torch.cuda.get_device_capability(0) == (9, 0):
+                assert self.config["self_attn_1_type"] != "sage_attn2"
             if "offload_ratio" in self.config:
                 offload_ratio = self.config["offload_ratio"]
             else: