"docs/vscode:/vscode.git/clone" did not exist on "a42736bbb8feeb047c240c062a0ba9de03c2ab40"
saint.cpp 624 Bytes
Newer Older
rusty1s's avatar
rusty1s committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
#include <Python.h>
#include <torch/script.h>

#include "cpu/saint_cpu.h"

#ifdef _WIN32
PyMODINIT_FUNC PyInit__saint(void) { return NULL; }
#endif

std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>
subgraph(torch::Tensor idx, torch::Tensor rowptr, torch::Tensor row,
         torch::Tensor col) {
  if (idx.device().is_cuda()) {
#ifdef WITH_CUDA
    AT_ERROR("No CUDA version supported");
#else
    AT_ERROR("Not compiled with CUDA support");
#endif
  } else {
    return subgraph_cpu(idx, rowptr, row, col);
  }
}

static auto registry =
    torch::RegisterOperators().op("torch_sparse::saint_subgraph", &subgraph);