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
32
33
34
|
// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is licensed under the BSD-style license found in the
// LICENSE file in the root directory of this source tree.
#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))
|
