[GPT OSS] Fix false flag (#43120)

fix

Author: vasqu

src/transformers/models/gpt_oss/modeling_gpt_oss.py

@@ -434,7 +434,7 @@ class GptOssPreTrainedModel(PreTrainedModel):
     _skip_keys_device_placement = ["past_key_values"]
     _supports_flash_attn = True
     _supports_sdpa = False
-    _supports_flex_attn = False
+    _supports_flex_attn = True
 
     _can_compile_fullgraph = True
     _supports_attention_backend = True

src/transformers/models/gpt_oss/modular_gpt_oss.py

@@ -354,7 +354,6 @@ def forward(
 class GptOssPreTrainedModel(LlamaPreTrainedModel):
     _keep_in_fp32_modules = ["post_attention_layernorm", "input_layernorm", "norm"]
     _supports_sdpa = False
-    _supports_flex_attn = False
     _can_record_outputs = {
         "router_logits": OutputRecorder(GptOssTopKRouter, index=0),
         "hidden_states": GptOssDecoderLayer,