pytorch/test/cpp/api/tensor_cuda.cpp

12 lines
299 B
C++
Raw Normal View History

#include <gtest/gtest.h>
2018-06-20 18:44:21 +00:00
#include <ATen/ATen.h>
#include <cmath>
TEST(Tensor, AllocatesTensorOnTheCorrectDevice_MultiCUDA) {
2018-06-20 18:44:21 +00:00
auto tensor = at::tensor({1, 2, 3}, at::device({at::kCUDA, 1}));
ASSERT_EQ(tensor.device().type(), at::Device::Type::CUDA);
ASSERT_EQ(tensor.device().index(), 1);
2018-06-20 18:44:21 +00:00
}