Autotuning cache miss for PallasTritonGatedLinearUnit and Inference time increases significantly.

After updating to the latest version of the code, I noticed the following Warning when running on an RTX A6000: `Autotuning cache miss for PallasTritonGatedLinearUnit(config=None, vjp=GatedLinearUnitVjp(config=None, vjp=None)) on NVIDIA RTX A6000 with key immutabledict.` Additionally, for a protein with a length of 898 amino acids, the GPU inference time for a single seed is at least 80 seconds longer compared to version V3.0.1.