diff --git a/src/core/shape_inference/include/one_hot_shape_inference.hpp b/src/core/shape_inference/include/one_hot_shape_inference.hpp index f6f944de456680..ece6634cf638bc 100644 --- a/src/core/shape_inference/include/one_hot_shape_inference.hpp +++ b/src/core/shape_inference/include/one_hot_shape_inference.hpp @@ -67,7 +67,7 @@ std::vector shape_infer(const OneHot* op, auto depth_as_shape = get_input_const_data_as_shape(op, 1, ta, util::GetNotNegative(op)); - if (depth_as_shape && depth_as_shape->size() == 1 && (*depth_as_shape)[0].get_length() > 0) { + if (depth_as_shape && depth_as_shape->size() == 1) { result_shape.insert(result_shape.begin() + axis, (*depth_as_shape)[0]); } else { result_shape.insert(result_shape.begin() + axis, DimType()); diff --git a/src/plugins/intel_gpu/include/intel_gpu/primitives/one_hot.hpp b/src/plugins/intel_gpu/include/intel_gpu/primitives/one_hot.hpp index caec53c857378c..d47a128c62c434 100644 --- a/src/plugins/intel_gpu/include/intel_gpu/primitives/one_hot.hpp +++ b/src/plugins/intel_gpu/include/intel_gpu/primitives/one_hot.hpp @@ -63,7 +63,7 @@ struct one_hot : public primitive_base { const int64_t& one_hot_axis, const float& on_value = 1.0f, const float& off_value = 0.0f) - : primitive_base(id, {input, input_depth}) + : primitive_base(id, {input, input_depth}, 1, {optional_data_type{output_dt}}) , shape(shape) , one_hot_axis(one_hot_axis) , on_value(on_value) diff --git a/src/plugins/intel_gpu/src/graph/one_hot.cpp b/src/plugins/intel_gpu/src/graph/one_hot.cpp index fe810a9bd8bf5c..e9613f360fd380 100644 --- a/src/plugins/intel_gpu/src/graph/one_hot.cpp +++ b/src/plugins/intel_gpu/src/graph/one_hot.cpp @@ -66,14 +66,15 @@ std::vector one_hot_inst::calc_output_layouts(const one_hot_node& /*node }; int64_t depth = desc->depth; - - auto depth_tensor = ov::Tensor(ov::element::i64, ov::Shape{1}, static_cast(&depth)); - std::unordered_map const_data = { - {1, depth_tensor} - }; - auto& memory_deps = impl_param.memory_deps; - if (memory_deps.count(1) > 0) { + + std::unordered_map const_data = {}; + if (depth != 0) { + auto depth_tensor = ov::Tensor(ov::element::i64, ov::Shape{1}, static_cast(&depth)); + const_data = { + {1, depth_tensor} + }; + } else if (memory_deps.count(1) > 0) { auto depth_mem = memory_deps.at(1); cldnn::mem_lock depth_lock(depth_mem, impl_param.get_stream());