remove torch.jit

This decorator will fail in some cases. Remove it until we figure out a solution.

modeling_flash_llama.py

@@ -62,7 +62,7 @@ logger = logging.get_logger(__name__)
 _CONFIG_FOR_DOC = "LlamaConfig"
 
 
-@torch.jit.script
+# @torch.jit.script
 def rmsnorm_func(hidden_states, weight, variance_epsilon):
     input_dtype = hidden_states.dtype
     hidden_states = hidden_states.to(torch.float32)