mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
1dd503c6fb
Fixes #ISSUE_NUMBER Pull Request resolved: https://github.com/pytorch/pytorch/pull/139256 Approved by: https://github.com/ezyang
25 lines
658 B
C++
25 lines
658 B
C++
// Copyright 2004-present Facebook. All Rights Reserved.
|
|
#define TORCH_ASSERT_ONLY_METHOD_OPERATORS
|
|
#include <ATen/core/Tensor.h>
|
|
|
|
#include <ATen/Dispatch.h>
|
|
#include <ATen/native/cpu/StackKernel.h>
|
|
#include <ATen/native/cpu/SerialStackImpl.h>
|
|
|
|
namespace at::native {
|
|
|
|
namespace {
|
|
|
|
void stack_serial_kernel(Tensor& result, TensorList tensors, int64_t dim) {
|
|
AT_DISPATCH_FLOATING_TYPES(
|
|
result.scalar_type(), "stack_serial_kernel", [&]() {
|
|
detail::stack_serial_kernel_impl<scalar_t, TensorList>(result, tensors, dim);
|
|
});
|
|
}
|
|
|
|
} // anonymous namespace
|
|
|
|
REGISTER_DISPATCH(stack_serial_stub, &stack_serial_kernel)
|
|
|
|
} // namespace at::native
|