"vscode:/vscode.git/clone" did not exist on "82febc59ad1d068144ee6d9fba7766628b521bbb"
graclus.cpp 329 Bytes
Newer Older
rusty1s's avatar
new try  
rusty1s committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
#include <torch/torch.h>

#include "../include/degree.cpp"
#include "../include/loop.cpp"

at::Tensor graclus(at::Tensor row, at::Tensor col, int num_nodes) {
  CHECK_CUDA(row);
  CHECK_CUDA(col);

  std::tie(row, col) = remove_self_loops(row, col);
  auto deg = degree(row, num_nodes, row.type().scalarType());

  return deg;
}