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
26
27
28
29
30
31
|
// (c) Facebook, Inc. and its affiliates. Confidential and proprietary.
#pragma once
#include <ATen/ATen.h>
#include <typeinfo>
inline std::string torch_tensor_device_name(const at::Tensor& ten) {
return c10::DeviceTypeName(ten.device().type());
}
#define TENSOR_NDIM_EQUALS(ten, dims) \
TORCH_CHECK( \
(ten).ndimension() == (dims), \
"Tensor '" #ten "' must have " #dims \
" dimension(s). " \
"Found ", \
(ten).ndimension())
#define TENSOR_ON_CPU(x) \
TORCH_CHECK( \
!x.is_cuda(), \
#x " must be a CPU tensor; it is currently on device ", \
torch_tensor_device_name(x))
#define TENSOR_ON_CUDA_GPU(x) \
TORCH_CHECK( \
x.is_cuda(), \
#x " must be a CUDA tensor; it is currently on device ", \
torch_tensor_device_name(x))
|
