diff --git a/pystencils/gpucuda/texture_utils.py b/pystencils/gpucuda/texture_utils.py
index ff3db430fc841f4121b9602b57a1c1aa08ef86f9..8d5ba56c8ea056b38a5e21869fbd3ac8d89896e2 100644
--- a/pystencils/gpucuda/texture_utils.py
+++ b/pystencils/gpucuda/texture_utils.py
@@ -9,6 +9,7 @@
 """
 
 from os.path import dirname, isdir, join
+from typing import Union
 
 import numpy as np
 
@@ -29,7 +30,7 @@ def pow_two_divider(n):
     return divider
 
 
-def ndarray_to_tex(tex_ref,
+def ndarray_to_tex(tex_ref,  # type: Union[cuda.TextureReference, cuda.SurfaceReference]
                    ndarray,
                    address_mode=None,
                    filter_mode=None,
@@ -51,7 +52,7 @@ def ndarray_to_tex(tex_ref,
         raise TypeError(
             'ndarray must be numpy.ndarray or pycuda.gpuarray.GPUArray')
 
-    cuda.TextureReference.set_array(tex_ref, cu_array)
+    tex_ref.set_array(cu_array)
 
     tex_ref.set_address_mode(0, address_mode)
     if ndarray.ndim >= 2: