Update

flux/math.py

@@ -11,7 +11,7 @@ def check_tpu():
     return any('TPU' in d.device_kind for d in jax.devices())
 
 # from torch import Tensor
-if check_tpu():
+if False:
     from jax.experimental.pallas.ops.tpu.flash_attention import flash_attention
     # q,  # [batch_size, num_heads, q_seq_len, d_model]
     # k,  # [batch_size, num_heads, kv_seq_len, d_model]