From 02bbb215873549a2e999cfc8a635097d904932ad Mon Sep 17 00:00:00 2001
From: Stephan Seitz <stephan.seitz@fau.de>
Date: Thu, 24 Oct 2019 16:36:55 +0200
Subject: [PATCH] Give every each torch operation its own build_directory
---
src/pystencils_autodiff/backends/astnodes.py | 8 ++++++--
1 file changed, 6 insertions(+), 2 deletions(-)
diff --git a/src/pystencils_autodiff/backends/astnodes.py b/src/pystencils_autodiff/backends/astnodes.py
index 0a1e5c0..acd536f 100644
--- a/src/pystencils_autodiff/backends/astnodes.py
+++ b/src/pystencils_autodiff/backends/astnodes.py
@@ -8,6 +8,7 @@
"""
+import os
import sys
from collections.abc import Iterable
from os.path import dirname, exists, join
@@ -102,12 +103,15 @@ class TorchModule(JinjaCppFile):
if not exists(file_name):
write_file(file_name, source_code)
- # TODO: propagate extra headers
+
+ build_dir = join(get_cache_config()['object_cache'], self.module_name)
+ os.makedirs(build_dir, exist_ok=True)
+
torch_extension = load(hash,
[file_name],
with_cuda=self.is_cuda,
extra_cflags=['--std=c++14'],
- build_directory=get_cache_config()['object_cache'],
+ build_directory=build_dir,
extra_include_paths=[get_pycuda_include_path(),
get_pystencils_include_path()])
return torch_extension
--
GitLab