Just wanted to add my +1 on this - getting this properly resolved would be really helpful for us (we have several J401s, and the JP5 CUDA version is barring us from upgrading to the latest torch/JAX). Unfortunately we haven’t had any luck with the workarounds mentioned above.