Update acc_kernel.cpp

Oneflow-Inc · Dec 20, 2023 · 20cc5c3 · 20cc5c3
1 parent 97292e2
commit 20cc5c3
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/oneflow/user/kernels/acc_kernel.cpp b/oneflow/user/kernels/acc_kernel.cpp
@@ -36,6 +36,7 @@ class AccKernel final : public user_op::OpKernel {
     std::unique_ptr<ep::primitive::Add> primitive =
         ep::primitive::NewPrimitive<ep::primitive::AddFactory>(ctx->device_type(), in->data_type());
     CHECK(primitive);
+    std::cout << ctx->Tensor4ArgNameAndIndex("in", 0)->shape_view().elem_cnt() << std::endl;
     primitive->Launch(ctx->stream(), out->dptr(), in->dptr(), out->mut_dptr(),
                       in->shape_view().elem_cnt());
   }