mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 00:21:07 +01:00
3ee97183b0
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/19660 Implementation of aggregated Scale operator. The operator takes a list of tensors as an input and scales all of them them with the argument float value. The tensor sizes can be different, therefore bookkeeping of the sizes and pointers to the tensors are necessary for the GPU version of the kernel. Reviewed By: BIT-silence Differential Revision: D14984233 fbshipit-source-id: 37cc97159a4f2c38cd6fff4f5710ab7d3a773611
19 lines
553 B
C++
19 lines
553 B
C++
#include "caffe2/operators/scale_blobs_op.h"
|
|
|
|
namespace caffe2 {
|
|
|
|
REGISTER_CPU_OPERATOR(ScaleBlobs, ScaleBlobsOp<CPUContext>);
|
|
OPERATOR_SCHEMA(ScaleBlobs)
|
|
.NumInputs(1, INT_MAX)
|
|
.NumOutputs(1, INT_MAX)
|
|
.AllowInplace([](int, int) { return true; })
|
|
.IdenticalTypeAndShape()
|
|
.SetDoc(R"DOC(
|
|
ScaleBlobs takes one or more input data (Tensor) and produces one
|
|
or more output data (Tensor) whose value is the input data tensor
|
|
scaled element-wise.
|
|
)DOC")
|
|
.Arg("scale", "(float, default 1.0) the scale to apply.");
|
|
|
|
} // namespace caffe2
|