bring back TINYGRAD_DEBUG. not sure why it was removed

exo-explore · Jul 31, 2024 · 980d5d2 · 980d5d2
1 parent 7676625
commit 980d5d2
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/exo/helpers.py b/exo/helpers.py
@@ -36,6 +36,8 @@ def get_inference_engine(inference_engine_name):
     return MLXDynamicShardInferenceEngine()
   elif inference_engine_name == "tinygrad":
     from exo.inference.tinygrad.inference import TinygradDynamicShardInferenceEngine
+    import tinygrad.helpers
+    tinygrad.helpers.DEBUG.value = int(os.getenv("TINYGRAD_DEBUG", default="0"))
 
     return TinygradDynamicShardInferenceEngine()
   else: