/usr/local/lib64/python3.6/site-packages/torch/include/ATen/cuda
Edit: /usr/local/lib64/python3.6/site-packages/torch/include/ATen/cuda/CUDAUtils.h (428B)
#pragma once
#include
namespace at { namespace cuda {
// Check if every tensor in a list of tensors matches the current
// device.
inline bool check_device(ArrayRef ts) {
if (ts.empty()) {
return true;
}
Device curDevice = Device(kCUDA, current_device());
for (const Tensor& t : ts) {
if (t.device() != curDevice) return false;
}
return true;
}
}} // namespace at::cuda