add assert torch_ring_attention is not None

chengzeyi · Dec 25, 2024 · ee9e03b · ee9e03b
1 parent fbb42ff
commit ee9e03b
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/src/para_attn/para_attn_interface.py b/src/para_attn/para_attn_interface.py
@@ -119,6 +119,7 @@ def forward(
         mesh,
     ):
         assert _templated_ring_attention is not None, "RingAttnFunc requires a newer version of PyTorch"
+        assert torch_ring_attention is not None, "RingAttnFunc requires a newer version of PyTorch"
 
         with unittest.mock.patch.object(
             torch_ring_attention,