diff --git a/src/pystencils_autodiff/backends/_torch_native.py b/src/pystencils_autodiff/backends/_torch_native.py index d63db1dd52bed945fb31cc34c8ff6027061d62c2..53619b5e2504b5926bcd2cf6defd026a5af14c73 100644 --- a/src/pystencils_autodiff/backends/_torch_native.py +++ b/src/pystencils_autodiff/backends/_torch_native.py @@ -106,8 +106,7 @@ def generate_torch(destination_folder, } if is_cuda: - template_string_cpp = read_file(join(dirname(__file__), - 'torch_native_cuda.tmpl.cpp')) + template_string_cpp = read_file(join(dirname(__file__), 'torch_native_cuda.tmpl.cpp')) template = jinja2.Template(template_string_cpp) output = template.render(render_dict) write_file(join(destination_folder, operation_string + '.cpp'), output) @@ -117,8 +116,7 @@ def generate_torch(destination_folder, output = template.render(render_dict) write_file(join(destination_folder, operation_string + '.cu'), output) else: - template_string_cpp = read_file(join(dirname(__file__), - 'torch_native_cpu.tmpl.cpp')) + template_string_cpp = read_file(join(dirname(__file__), 'torch_native_cpu.tmpl.cpp')) template = jinja2.Template(template_string_cpp) output = template.render(render_dict) write_file(join(destination_folder, operation_string + '.cpp'), output) diff --git a/src/pystencils_autodiff/backends/python_bindings.py b/src/pystencils_autodiff/backends/python_bindings.py index 3cae1a83d7f6b09203018aab4452a807b5a074ef..d7acce1c79bbc8c5e11b3642da5b290705bbff30 100644 --- a/src/pystencils_autodiff/backends/python_bindings.py +++ b/src/pystencils_autodiff/backends/python_bindings.py @@ -12,7 +12,7 @@ import itertools import jinja2 import stringcase -from pystencils.astnodes import KernelFunction, Node +from pystencils.astnodes import KernelFunction from pystencils_autodiff.framework_integration.astnodes import JinjaCppFile diff --git a/src/pystencils_autodiff/tensorflow_jit.py b/src/pystencils_autodiff/tensorflow_jit.py index 76e887c83264c36f9cbd8ec17cf568fe406c50d6..be3fbf98ebe2c5eea30f850a59415e8d1f6c1b07 100644 --- a/src/pystencils_autodiff/tensorflow_jit.py +++ b/src/pystencils_autodiff/tensorflow_jit.py @@ -10,7 +10,6 @@ import json import subprocess import sysconfig -from itertools import chain from os.path import exists, join from tqdm import tqdm