Update functions to look in chroma/cuda directory for CUDA source

1 files changed, 5 insertions, 5 deletions

diff --git a/chroma/gpu.py b/chroma/gpu.py
index 4e2ecf3..2b143a3 100644
--- a/chroma/gpu.py
+++ b/chroma/gpu.py
@@ -13,7 +13,7 @@ from pycuda import characterize
 import pycuda.driver as cuda
 from pycuda import gpuarray as ga
 
-import chroma.src
+import chroma.cuda
 from chroma.tools import timeit, profile_if_possible
 from chroma.geometry import standard_wavelengths
 from chroma import event
@@ -26,7 +26,7 @@ cuda_options = ('--use_fast_math',)#, '--ptxas-options=-v']
 @pycuda.tools.context_dependent_memoize
 def get_cu_module(name, options=None, include_source_directory=True):
     """Returns a pycuda.compiler.SourceModule object from a CUDA source file
-    located in the chroma src directory at src/[name]."""
+    located in the chroma cuda directory at cuda/[name]."""
     if options is None:
         options = []
     elif isinstance(options, tuple):
@@ -34,7 +34,7 @@ def get_cu_module(name, options=None, include_source_directory=True):
     else:
         raise TypeError('`options` must be a tuple.')
 
-    srcdir = os.path.dirname(os.path.abspath(chroma.src.__file__))
+    srcdir = os.path.dirname(os.path.abspath(chroma.cuda.__file__))
 
     if include_source_directory:
         options += ['-I' + srcdir]
@@ -47,9 +47,9 @@ def get_cu_module(name, options=None, include_source_directory=True):
 
 @pycuda.tools.memoize
 def get_cu_source(name):
-    """Get the source code for a CUDA source file located in the chroma src
+    """Get the source code for a CUDA source file located in the chroma cuda
     directory at src/[name]."""
-    srcdir = os.path.dirname(os.path.abspath(chroma.src.__file__))
+    srcdir = os.path.dirname(os.path.abspath(chroma.cuda.__file__))
     with open('%s/%s' % (srcdir, name)) as f:
         source = f.read()
     return source