In alphafold3, jax can't prof gpu data， only can prof the cpu and memcpy data.

[south-ocean]

In alphafold3, i am append code as belows:

with above prof code, i can't get the device data and No capturing function was found! in cupti_tracer.cc:

It can only get the cpu and memcpy data, but i can;t see the gpu kernel, So does this normal or what should i do? thannk.

[south-ocean]

@jakevdp Please help check this issue, thanks.