add PARROTS_JIT_OPTION (#816)

* code update * add PARROTS_JIT_OPTION * rm flow_warp_module * fix for lint * small fix
open-mmlab · Feb 5, 2021 · 4a694cc · 4a694cc
1 parent 5df759a
commit 4a694cc
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/mmcv/utils/parrots_jit.py b/mmcv/utils/parrots_jit.py
@@ -1,6 +1,10 @@
+import os
+
 from .parrots_wrapper import TORCH_VERSION
 
-if TORCH_VERSION == 'parrots':
+parrots_jit_option = os.getenv('PARROTS_JIT_OPTION')
+
+if TORCH_VERSION == 'parrots' and parrots_jit_option == 'ON':
     from parrots.jit import pat as jit
 else: