ScriptModuleWrapper in a sepatate file

and imported where needed
dbolya · Feb 12, 2020 · fccd89b · fccd89b
1 parent cf2d4e6
commit fccd89b
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 18 deletions.
diff --git a/layers/modules/fast_mask_iou.py b/layers/modules/fast_mask_iou.py
@@ -5,16 +5,7 @@
 #locals
 from data.config import Config
 from utils.functions import make_net
-
-
-#  As of March 10, 2019, Pytorch DataParallel still doesn't support JIT Script Modules
-use_jit = torch.cuda.device_count() <= 1
-if not use_jit:
-    print('Multiple GPUs detected! Turning off JIT.')
-
-ScriptModuleWrapper = torch.jit.ScriptModule if use_jit else nn.Module
-script_method_wrapper = torch.jit.script_method if use_jit else lambda fn, _rcn=None: fn
-
+from utils.script_module_wrapper import ScriptModuleWrapper, script_method_wrapper
 
 class FastMaskIoUNet(ScriptModuleWrapper):
 

diff --git a/layers/modules/fpn.py b/layers/modules/fpn.py
@@ -7,14 +7,7 @@
 
 #local imports
 from data.config import Config
-
-#  As of March 10, 2019, Pytorch DataParallel still doesn't support JIT Script Modules
-use_jit = torch.cuda.device_count() <= 1
-if not use_jit:
-    print('Multiple GPUs detected! Turning off JIT.')
-
-ScriptModuleWrapper = torch.jit.ScriptModule if use_jit else nn.Module
-script_method_wrapper = torch.jit.script_method if use_jit else lambda fn, _rcn=None: fn
+from utils.script_module_wrapper import ScriptModuleWrapper, script_method_wrapper
 
 
 

diff --git a/utils/script_module_wrapper.py b/utils/script_module_wrapper.py
@@ -0,0 +1,10 @@
+import torch
+import torch.nn
+
+#  As of March 10, 2019, Pytorch DataParallel still doesn't support JIT Script Modules
+use_jit = torch.cuda.device_count() <= 1
+if not use_jit:
+    print('Multiple GPUs detected! Turning off JIT.')
+
+ScriptModuleWrapper = torch.jit.ScriptModule if use_jit else nn.Module #TODO remove once nn.Module supports JIT script modules
+script_method_wrapper = torch.jit.script_method if use_jit else lambda fn, _rcn=None: fn