提交 c58f2555 编写于 作者: A Abhinav Arora

Fix one_hot_op.h

上级 0e06e606
......@@ -58,7 +58,8 @@ class OneHotKernel : public framework::OpKernel<T> {
int depth = context.Attr<int>("depth");
framework::VisitDataType(
static_cast<framework::proto::DataType>(context.Attr<int>("dtype")),
static_cast<framework::proto::VarType::Type>(
context.Attr<int>("dtype")),
OneHotOpFunctor<DeviceContext, T>(
in, out, depth, context.template device_context<DeviceContext>()));
}
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册