mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
508f676c50
Summary: Codemod generated with clangr shard mode, 50 files per diff, clangr code(ndim()->dim()): diffusion/FBS/browse/master/fbcode/caffe2/caffe2/fb/codemods/TensorMethodRename.cpp Reviewed By: salexspb Differential Revision: D12935787 fbshipit-source-id: 303d71d3eb050789af2ab9575e5dcc48f6037086
53 lines
1.4 KiB
C++
53 lines
1.4 KiB
C++
#include <iostream>
|
|
|
|
#include <gtest/gtest.h>
|
|
#include "caffe2/core/context.h"
|
|
#include "caffe2/core/context_gpu.h"
|
|
#include "caffe2/core/flags.h"
|
|
#include "caffe2/operators/reshape_op.h"
|
|
#include "caffe2/utils/math.h"
|
|
|
|
C10_DECLARE_string(caffe_test_root);
|
|
|
|
namespace caffe2 {
|
|
|
|
static void AddConstInput(
|
|
const vector<int64_t>& shape,
|
|
const float value,
|
|
const string& name,
|
|
Workspace* ws) {
|
|
DeviceOption option;
|
|
option.set_device_type(PROTO_CUDA);
|
|
CUDAContext context(option);
|
|
Blob* blob = ws->CreateBlob(name);
|
|
auto* tensor = BlobGetMutableTensor(blob, CUDA);
|
|
tensor->Resize(shape);
|
|
math::Set<float, CUDAContext>(
|
|
tensor->numel(), value, tensor->template mutable_data<float>(), &context);
|
|
return;
|
|
}
|
|
|
|
TEST(ReshapeOpGPUTest, testReshapeWithScalar) {
|
|
if (!HasCudaGPU())
|
|
return;
|
|
Workspace ws;
|
|
OperatorDef def;
|
|
def.set_name("test_reshape");
|
|
def.set_type("Reshape");
|
|
def.add_input("X");
|
|
def.add_output("XNew");
|
|
def.add_output("OldShape");
|
|
def.add_arg()->CopyFrom(MakeArgument("shape", vector<int64_t>{1}));
|
|
def.mutable_device_option()->set_device_type(PROTO_CUDA);
|
|
AddConstInput(vector<int64_t>(), 3.14, "X", &ws);
|
|
// execute the op
|
|
unique_ptr<OperatorBase> op(CreateOperator(def, &ws));
|
|
EXPECT_TRUE(op->Run());
|
|
Blob* XNew = ws.GetBlob("XNew");
|
|
const Tensor& XNewTensor = XNew->Get<Tensor>();
|
|
EXPECT_EQ(1, XNewTensor.dim());
|
|
EXPECT_EQ(1, XNewTensor.numel());
|
|
}
|
|
|
|
} // namespace caffe2
|