From 4ac1503bd507465edc6e58e9025c8a22867a0d20 Mon Sep 17 00:00:00 2001 From: Cjkkkk Date: Thu, 14 Mar 2024 14:24:29 -0700 Subject: [PATCH] rm XLA_FLAGS --- tests/fused_attention_stablehlo_test.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/fused_attention_stablehlo_test.py b/tests/fused_attention_stablehlo_test.py index a5be69ae6..ed76e2f52 100644 --- a/tests/fused_attention_stablehlo_test.py +++ b/tests/fused_attention_stablehlo_test.py @@ -16,7 +16,7 @@ from functools import partial from absl.testing import absltest from typing import Optional import os -os.environ['XLA_FLAGS'] = '--xla_gpu_enable_cudnn_fmha=true --xla_gpu_fused_attention_use_cudnn_rng=true --xla_dump_hlo_as_text --xla_dump_to=./hlo' +os.environ['XLA_FLAGS'] = '--xla_gpu_enable_cudnn_fmha=true --xla_gpu_fused_attention_use_cudnn_rng=true' import numpy as np import jax