Add JIT_OPTIMIZE env var

hydra_node/config.py

@@ -149,6 +149,7 @@ def init_config_model():
     config.default_config = os.getenv('DEFAULT_CONFIG', None)
     config.quality_hack = os.getenv('QUALITY_HACK', "0")
     config.clip_contexts = os.getenv('CLIP_CONTEXTS', "1")
+    config.jit_optimize = os.getenv('JIT_OPTIMIZE', "0")
     try:
         config.clip_contexts = int(config.clip_contexts)
         if config.clip_contexts < 1 or config.clip_contexts > 10:
@@ -189,6 +190,10 @@ def init_config_model():
             modules = load_modules(config.module_path)
             #attach it to the model
             model.premodules = modules
+
+        # enable JIT
+        if config.jit_optimize == "1":
+            model.fuse_model(ema=config.enable_ema == "1")
     
     config.model = model