TORCHINDUCTOR DISABLE CACHES

pytorch · Oct 31, 2024 · d03db9f · d03db9f
1 parent 3594697
commit d03db9f
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/tests/recipes/dev/test_generate_v2.py b/tests/recipes/dev/test_generate_v2.py
@@ -77,13 +77,12 @@ def test_llama2_generate_with_quantization(self, caplog, monkeypatch, tmpdir):
         tokenizer_path = Path(TOKENIZER_PATHS["llama2"])
         ckpt_dir = ckpt_path.parent
 
-        torch._dynamo.reset()
-
         # Config file needed for model conversion.
         write_hf_ckpt_config(ckpt_dir)
 
         cmd = f"""
-        tune run dev/generate_v2 \
+        TORCHINDUCTOR_FORCE_DISABLE_CACHES=1 \
+            tune run dev/generate_v2 \
             --config llama2/generation_v2 \
             output_dir={tmpdir} \
             checkpointer=torchtune.training.FullModelTorchTuneCheckpointer \