PixelShuffleKernel.h 334 B

1234567891011121314
  1. #pragma once
  2. #include <ATen/native/DispatchStub.h>
  3. namespace at {
  4. class TensorBase;
  5. }
  6. namespace at { namespace native {
  7. using pixel_shuffle_fn = void(*)(TensorBase&, const TensorBase&, int64_t);
  8. DECLARE_DISPATCH(pixel_shuffle_fn, pixel_shuffle_kernel);
  9. DECLARE_DISPATCH(pixel_shuffle_fn, pixel_unshuffle_kernel);
  10. }} // at::native