Error of fused_dense module of flash_attn

Hi author, I have compiled and installed fused_dense_lib successfully. But when I tried to run the finetuning code, I encountered this error: "RuntimeError: linear_act_forward failed." which is due to the line 291 in InternVideo2/single_modality/models/internvideo2.py. The complete error is as shown below:

Traceback (most recent call last):
File "", line 1, in
File "/nvme/miniconda3/envs/internvideo/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
return forward_call(*input, **kwargs)
File "/nvme/miniconda3/envs/internvideo/lib/python3.8/site-packages/flash_attn/ops/fused_dense.py", line 457, in forward
out = fused_mlp_func(
File "/nvme/miniconda3/envs/internvideo/lib/python3.8/site-packages/flash_attn/ops/fused_dense.py", line 391, in fused_mlp_func
return FusedMLPFunc.apply(
File "/nvme/miniconda3/envs/internvideo/lib/python3.8/site-packages/torch/cuda/amp/autocast_mode.py", line 97, in decorate_fwd
return fwd(*args, **kwargs)
File "/nvme/miniconda3/envs/internvideo/lib/python3.8/site-packages/flash_attn/ops/fused_dense.py", line 257, in forward
output1, *rest = fused_dense_cuda.linear_act_forward(
RuntimeError: linear_act_forward failed.

Could you please help me to fix this? Thank you very much.

Start at line 291 in InternVideo2/single_modality/models/internvideo2.py and reproduce the finetuning run that reaches flash_attn/ops/fused_dense.py. Trace the call through fused_mlp_func to linear_act_forward, then identify the environment or input condition causing the failure; done means the finetuning path runs without this RuntimeError.