pytorch/test/cpp/api/enum.cpp
Will Feng aad5071206 Use torch::variant for enums in C++ API
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/26837

Test Plan: Imported from OSS

Differential Revision: D17579438

Pulled By: yf225

fbshipit-source-id: 9ac59df28a317fdb3be2cc02c65962ad99117127
2019-10-16 22:40:57 -07:00

44 lines
1.2 KiB
C++

#include <gtest/gtest.h>
#include <c10/util/variant.h>
#include <torch/torch.h>
#include <test/cpp/api/support.h>
#define TORCH_ENUM_PRETTY_PRINT_TEST(name) \
{ \
v = torch::k##name; \
ASSERT_EQ(c10::visit(torch::enumtype::enum_name{}, v), #name); \
}
TEST(EnumTest, AllEnums) {
c10::variant<
torch::enumtype::kLinear,
torch::enumtype::kConv1D,
torch::enumtype::kConv2D,
torch::enumtype::kConv3D,
torch::enumtype::kConvTranspose1D,
torch::enumtype::kConvTranspose2D,
torch::enumtype::kConvTranspose3D,
torch::enumtype::kSigmoid,
torch::enumtype::kTanh,
torch::enumtype::kReLU,
torch::enumtype::kLeakyReLU,
torch::enumtype::kFanIn,
torch::enumtype::kFanOut
> v;
TORCH_ENUM_PRETTY_PRINT_TEST(Linear)
TORCH_ENUM_PRETTY_PRINT_TEST(Conv1D)
TORCH_ENUM_PRETTY_PRINT_TEST(Conv2D)
TORCH_ENUM_PRETTY_PRINT_TEST(Conv3D)
TORCH_ENUM_PRETTY_PRINT_TEST(ConvTranspose1D)
TORCH_ENUM_PRETTY_PRINT_TEST(ConvTranspose2D)
TORCH_ENUM_PRETTY_PRINT_TEST(ConvTranspose3D)
TORCH_ENUM_PRETTY_PRINT_TEST(Sigmoid)
TORCH_ENUM_PRETTY_PRINT_TEST(Tanh)
TORCH_ENUM_PRETTY_PRINT_TEST(ReLU)
TORCH_ENUM_PRETTY_PRINT_TEST(LeakyReLU)
TORCH_ENUM_PRETTY_PRINT_TEST(FanIn)
TORCH_ENUM_PRETTY_PRINT_TEST(FanOut)
}