mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
ad0ef7ae48
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/15723 As title says. Reviewed By: jianyuh Differential Revision: D13578604 fbshipit-source-id: 0da0ac31ae83c1e0daa9077e878feb4deffed6a3
29 lines
671 B
C++
29 lines
671 B
C++
#pragma once
|
|
|
|
#include "caffe2/operators/copy_op.h"
|
|
#include "caffe2/quantization/server/dnnlowp_op.h"
|
|
|
|
namespace caffe2 {
|
|
|
|
// FIXME
|
|
using BatchPermutationFP32Op = CopyOp<CPUContext, CPUContext, CPUContext>;
|
|
|
|
template <typename T>
|
|
class BatchPermutationDNNLowPOp final
|
|
: public DNNLowPOp<T, BatchPermutationFP32Op> {
|
|
public:
|
|
USE_OPERATOR_FUNCTIONS(CPUContext);
|
|
USE_DNNLOWP_OPERATOR_BASE_FUNCTIONS(T, BatchPermutationFP32Op);
|
|
|
|
BatchPermutationDNNLowPOp(const OperatorDef& operator_def, Workspace* ws)
|
|
: BaseType(operator_def, ws) {}
|
|
|
|
bool RunOnDevice() override;
|
|
|
|
private:
|
|
INPUT_TAGS(INPUT, INDICES);
|
|
OUTPUT_TAGS(OUTPUT);
|
|
};
|
|
|
|
} // namespace caffe2
|