Allow using nvFuser on CUDA extension (#1701)

I feel that it would be convenient to be able to run experiments and prototype something quickly in an extension rather than adding a new dummy test. Tested in https://github.com/zasdfgbnm/things/tree/master/2022/nvfuser-extension
pytorch · May 17, 2022 · 174d453 · 174d453
1 parent 18bee67
commit 174d453
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/setup.py b/setup.py
@@ -1047,6 +1047,9 @@ def print_box(msg):
                 'include/torch/csrc/jit/testing/*.h',
                 'include/torch/csrc/jit/tensorexpr/*.h',
                 'include/torch/csrc/jit/tensorexpr/operators/*.h',
+                'include/torch/csrc/jit/codegen/cuda/*.h',
+                'include/torch/csrc/jit/codegen/cuda/ops/*.h',
+                'include/torch/csrc/jit/codegen/cuda/scheduler/*.h',
                 'include/torch/csrc/onnx/*.h',
                 'include/torch/csrc/profiler/*.h',
                 'include/torch/csrc/utils/*.h',