diff options
Diffstat (limited to 'src/armnn/layers/Pooling2dLayer.hpp')
-rw-r--r-- | src/armnn/layers/Pooling2dLayer.hpp | 9 |
1 files changed, 5 insertions, 4 deletions
diff --git a/src/armnn/layers/Pooling2dLayer.hpp b/src/armnn/layers/Pooling2dLayer.hpp index af39dbb5ec..d5950d6ec3 100644 --- a/src/armnn/layers/Pooling2dLayer.hpp +++ b/src/armnn/layers/Pooling2dLayer.hpp @@ -9,19 +9,20 @@ namespace armnn { -class SoftmaxLayer : public LayerWithParameters<SoftmaxDescriptor> +class Pooling2dLayer : public LayerWithParameters<Pooling2dDescriptor> { public: virtual std::unique_ptr<IWorkload> CreateWorkload(const Graph& graph, const IWorkloadFactory& factory) const override; - SoftmaxLayer* Clone(Graph& graph) const override; + Pooling2dLayer* Clone(Graph& graph) const override; void ValidateTensorShapesFromInputs() override; + std::vector<TensorShape> InferOutputShapes(const std::vector<TensorShape>& inputShapes) const override; protected: - SoftmaxLayer(const SoftmaxDescriptor& param, const char* name); - ~SoftmaxLayer() = default; + Pooling2dLayer(const Pooling2dDescriptor& param, const char* name); + ~Pooling2dLayer() = default; }; } // namespace |