Hot fix after PyTorch pin update (#5017)

Signed-off-by: Anatoly Myachev <[email protected]>

Author: anmyachev

benchmarks/triton_kernels_benchmark/flex_attention_benchmark_causal_mask.py

@@ -13,7 +13,7 @@
 import torch._inductor.lowering
 import torch._inductor.kernel
 import torch._inductor.kernel.flex.flex_attention as flex_attn
-from torch._inductor.template_heuristics import FlexConfig, FlexDecodeConfig
+from torch._inductor.template_heuristics.triton import FlexConfig, FlexDecodeConfig
 
 import triton_kernels_benchmark as benchmark_suit
 import triton