mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-29 20:14:01 +00:00
29 lines
648 B
C++
29 lines
648 B
C++
// Copyright (c) Microsoft Corporation. All rights reserved.
|
|
// Licensed under the MIT License.
|
|
|
|
#pragma once
|
|
#include "core/common/common.h"
|
|
#include "core/framework/op_kernel.h"
|
|
#include "core/providers/cuda/cuda_common.h"
|
|
|
|
namespace onnxruntime {
|
|
namespace contrib {
|
|
namespace cuda {
|
|
|
|
using namespace onnxruntime::cuda;
|
|
|
|
template <typename T, typename U>
|
|
class LayerNorm final : public CudaKernel {
|
|
public:
|
|
LayerNorm(const OpKernelInfo& op_kernel_info);
|
|
|
|
Status ComputeInternal(OpKernelContext* ctx) const override;
|
|
|
|
private:
|
|
int64_t axis_;
|
|
double epsilon_;
|
|
};
|
|
|
|
} // namespace cuda
|
|
} // namespace contrib
|
|
} // namespace onnxruntime
|