#pragma once

#include "cpu/vision.h"

#ifdef WITH_CUDA

#include "cuda/vision.h"

#endif

// Interface for Python

at::Tensor GOF_forward(const at::Tensor& weight, 

                       const at::Tensor& gaborFilterBank) {

  if (weight.type().is_cuda()) {

#ifdef WITH_CUDA

    return GOF_forward_cuda(weight, gaborFilterBank);

#else

    AT_ERROR("Not compiled with GPU support");

#endif

  }

  return GOF_forward_cpu(weight, gaborFilterBank);

}

at::Tensor GOF_backward(const at::Tensor& grad_output,

                        const at::Tensor& gaborFilterBank) {

  if (grad_output.type().is_cuda()) {

#ifdef WITH_CUDA

    return GOF_backward_cuda(grad_output, gaborFilterBank);

#else

    AT_ERROR("Not compiled with GPU support");

#endif

  }

  return GOF_backward_cpu(grad_output, gaborFilterBank);

}