Set backend to eager

pytorch · Oct 31, 2024 · ae2cf5b · ae2cf5b
1 parent d03db9f
commit ae2cf5b
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/tests/recipes/dev/test_generate_v2.py b/tests/recipes/dev/test_generate_v2.py
@@ -81,8 +81,7 @@ def test_llama2_generate_with_quantization(self, caplog, monkeypatch, tmpdir):
         write_hf_ckpt_config(ckpt_dir)
 
         cmd = f"""
-        TORCHINDUCTOR_FORCE_DISABLE_CACHES=1 \
-            tune run dev/generate_v2 \
+        tune run dev/generate_v2 \
             --config llama2/generation_v2 \
             output_dir={tmpdir} \
             checkpointer=torchtune.training.FullModelTorchTuneCheckpointer \
@@ -101,6 +100,8 @@ def test_llama2_generate_with_quantization(self, caplog, monkeypatch, tmpdir):
         model_config = MODEL_TEST_CONFIGS["llama2"]
         cmd = cmd + model_config
 
+        import os
+        os.environ["TORCH_COMPILE_BACKEND"] = "eager"
         monkeypatch.setattr(sys, "argv", cmd)
         with pytest.raises(SystemExit, match=""):
             runpy.run_path(TUNE_PATH, run_name="__main__")