Encapsulate and standardize new_empty_tensor_op

torchvision/csrc/empty_tensor_op.h → torchvision/csrc/new_empty_tensor_op.cpp

@@ -1,29 +1,35 @@
 #pragma once
 
-// All pure C++ headers for the C++ frontend.
-#include <torch/all.h>
+#include "new_empty_tensor_op.h"
+#include <torch/extension.h>
+
+namespace {
 
 class NewEmptyTensorOp : public torch::autograd::Function<NewEmptyTensorOp> {
  public:
   static torch::autograd::variable_list forward(
       torch::autograd::AutogradContext* ctx,
-      torch::autograd::Variable input,
-      c10::List<int64_t> new_shape) {
+      const torch::autograd::Variable& input,
+      const c10::List<int64_t>& new_shape) {
     ctx->saved_data["shape"] = input.sizes();
     std::vector<int64_t> shape(new_shape.begin(), new_shape.end());
     return {input.new_empty(shape, at::TensorOptions())};
   }
 
   static torch::autograd::variable_list backward(
       torch::autograd::AutogradContext* ctx,
-      torch::autograd::variable_list grad_output) {
+      const torch::autograd::variable_list& grad_output) {
     // Use data saved in forward
     auto shape = ctx->saved_data["shape"].toIntList();
     auto out = forward(ctx, grad_output[0], shape);
     return {out[0], at::Tensor()};
   }
 };
 
-at::Tensor new_empty_tensor(const at::Tensor& input, c10::List<int64_t> shape) {
+} // namespace
+
+at::Tensor new_empty_tensor(
+    const at::Tensor& input,
+    const c10::List<int64_t>& shape) {
   return NewEmptyTensorOp::apply(input, shape)[0];
 }

torchvision/csrc/new_empty_tensor_op.h

@@ -0,0 +1,7 @@
+#pragma once
+
+#include <ATen/ATen.h>
+
+at::Tensor new_empty_tensor(
+    const at::Tensor& input,
+    const c10::List<int64_t>& shape);

torchvision/csrc/vision.cpp

@@ -9,7 +9,7 @@
 #endif
 
 #include "deform_conv2d.h"
-#include "empty_tensor_op.h"
+#include "new_empty_tensor_op.h"
 #include "nms.h"
 #include "ps_roi_align.h"
 #include "ps_roi_pool.h"