Utils.h 577 B

123456789101112131415161718192021
  1. #pragma once
  2. #include <ATen/core/Tensor.h>
  3. #include <ATen/cuda/Exceptions.h>
  4. #include <ATen/cudnn/cudnn-wrapper.h>
  5. #include <ATen/cudnn/Handle.h>
  6. namespace at { namespace native {
  7. // cuDNN has a buggy check for tensor being contiguous (that is, it does
  8. // not ignore stride for dimension that is equal to 0). This function
  9. // makes tensors which have zero stride contiguous, by setting the
  10. // strides to 1 as cuDNN likes.
  11. inline Tensor contiguousIfZeroInStrides(const Tensor& t) {
  12. for (auto s : t.strides()) {
  13. if (s == 0) return t.contiguous();
  14. }
  15. return t;
  16. }
  17. }}