Commit 0ef3d7a9 authored by Stephan Seitz's avatar Stephan Seitz
Browse files

Fix loading of tensorflow modules

parent d2c11529
......@@ -148,9 +148,9 @@ class TensorflowModule(TorchModule):
def compile(self):
from pystencils_autodiff.tensorflow_jit import compile_sources_and_load
if self.is_cuda:
self.compiled_file = compile_sources_and_load([], cuda_sources=[str(self)])
self.compiled_file = compile_sources_and_load([], cuda_sources=[str(self)], compile_only=True)
else:
self.compiled_file = compile_sources_and_load([str(self)])
self.compiled_file = compile_sources_and_load([str(self)], compile_only=True)
import tensorflow as tf
return tf.load_op_library(self.compiled_file)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment