cuda_utils.cpp 251 Bytes
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
#include <torch/extension.h>

int get_device_attribute(
    int attribute,
    int device_id);

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def(
    "get_device_attribute",
    &get_device_attribute,
    "Gets the specified device attribute.");
}