ChannelShuffleKernel.h 299 B

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