diff --git a/unsloth/kernels/flex_attention.py b/unsloth/kernels/flex_attention.py
index 7deb9a96d3..887ffca1b7 100644
--- a/unsloth/kernels/flex_attention.py
+++ b/unsloth/kernels/flex_attention.py
@@ -15,6 +15,7 @@
 import torch
 from functools import lru_cache
 from transformers.models.llama.modeling_llama import logger
+import os
 
 torch_compile_options = {
     "epilogue_fusion"   : True,