Check if inside a module or not

lorenzomammana · Apr 26, 2020 · aebd10c · aebd10c
1 parent a508ac9
commit aebd10c
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/morphology_cupy.py b/morphology_cupy.py
@@ -1,8 +1,12 @@
 import cupy as cp
 import numpy as np
 
-with open('morph_cupy/morphology.cu', 'r') as f:
-    code = f.read()
+try:
+    with open('morphology.cu', 'r') as f:
+        code = f.read()
+except FileNotFoundError:
+    with open('morph_cupy/morphology.cu', 'r') as f:
+        code = f.read()
 
 module = cp.RawModule(code=code)
 dilation_cuda = module.get_function('dilation')