From 18ed75f3fbba570293fda6e1d06030cdde916c3e Mon Sep 17 00:00:00 2001
From: Stephan Seitz <stephan.seitz@fau.de>
Date: Mon, 23 Sep 2019 14:37:30 +0200
Subject: [PATCH] Link against cudart
---
src/pystencils_autodiff/tensorflow_jit.py | 5 +++--
1 file changed, 3 insertions(+), 2 deletions(-)
diff --git a/src/pystencils_autodiff/tensorflow_jit.py b/src/pystencils_autodiff/tensorflow_jit.py
index 916912a..89984cc 100644
--- a/src/pystencils_autodiff/tensorflow_jit.py
+++ b/src/pystencils_autodiff/tensorflow_jit.py
@@ -14,9 +14,9 @@ import sysconfig
from os.path import exists, join
import p_tqdm
+
import pystencils
from pystencils.cpu.cpujit import get_cache_config, get_compiler_config, get_pystencils_include_path
-
from pystencils_autodiff._file_io import read_file, write_file
_hash = hashlib.md5
@@ -35,6 +35,7 @@ if get_compiler_config()['os'] != 'windows':
_position_independent_flag = "-fPIC"
_compile_env = os.environ.copy()
_object_file_extension = '.o'
+ _link_cudart = '-lcudart'
else:
_do_not_link_flag = '-c'
_output_flag = '-o'
@@ -75,7 +76,7 @@ def link(object_files, destination_file=None, overwrite_destination_file=True, a
*_tf_link_flags,
*_include_flags,
*additional_link_flags,
- '-lcudart',
+ _link_cudart,
_shared_object_flag,
_output_flag]
if not destination_file:
--
GitLab