logit_op.h

#include "caffe2/core/context.h"
#include "caffe2/core/operator.h"
#include "caffe2/utils/math.h"

namespace caffe2 {

template <typename T, class Context>
class LogitGradientOp final : public Operator<Context> {
 public:
  USE_OPERATOR_CONTEXT_FUNCTIONS;
  LogitGradientOp(const OperatorDef& operator_def, Workspace* ws)
      : Operator<Context>(operator_def, ws),
        eps_(OperatorBase::GetSingleArgument<float>("eps", 1e-6f)) {}
  ~LogitGradientOp() {}

  bool RunOnDevice() override;

 protected:
  float eps_;
};

} // namespace caffe2