diff --git a/src/pystencils_autodiff/backends/astnodes.py b/src/pystencils_autodiff/backends/astnodes.py index 10228c8b776ed9b157c3d0e2712b9235ffe2253b..a040b060adbe957e0e092e2ebc0617a47b63a576 100644 --- a/src/pystencils_autodiff/backends/astnodes.py +++ b/src/pystencils_autodiff/backends/astnodes.py @@ -154,7 +154,7 @@ class TorchModule(JinjaCppFile): os.makedirs(build_dir, exist_ok=True) file_name = join(build_dir, f'{hash}{file_extension}') - self.compiled_file = (join(build_dir, compile_module_name) or file_name).replace('.cpp', '') + '.so' + self.compiled_file = join(build_dir, compile_module_name or file_name).replace('.cpp', '') + '.so' if not exists(file_name): write_file(file_name, source_code)