Disable dynamo when running generated opcheck tests (pytorch#111685)

williamwen42 · xuhancn · commit 8393a2841ab4 · 2023-11-08T12:54:01.000Z
Summary: Use `TORCHDYNAMO_DISABLE=1` when running generated opcheck tests. Enable some `fbgemm::pack_segments` tests that errored out (with error `RuntimeError: expected int but got s0*s1**2`) because dynamo was being run in the opcheck tests. Test Plan: `parsh -v --build-flags mode/dev-nosan //deeplearning/fbgemm/fbgemm_gpu:sparse_ops_test` then `run_tests("test_pack_segments")` Differential Revision: D50508958 Pull Request resolved: pytorch#111685 Approved by: https://github.com/zou3519
diff --git a/torch/testing/_internal/optests/generate_tests.py b/torch/testing/_internal/optests/generate_tests.py
@@ -449,11 +449,14 @@ def maybe_raise_errors_on_exit(self) -> None:
 
     def __enter__(self, *args, **kwargs):
         self.prev_is_opcheck_mode = _is_inside_opcheck_mode.value
+        self.prev_dynamo_disable = os.environ.get("TORCHDYNAMO_DISABLE", "")
         _is_inside_opcheck_mode.value = True
+        os.environ["TORCHDYNAMO_DISABLE"] = "1"
         return super().__enter__(*args, **kwargs)
 
     def __exit__(self, *args, **kwargs):
         _is_inside_opcheck_mode.value = self.prev_is_opcheck_mode
+        os.environ["TORCHDYNAMO_DISABLE"] = self.prev_dynamo_disable
         try:
             self.maybe_raise_errors_on_exit()
             if should_update_failures_dict():
