Program Listing for File custom_op.h
↰ Return to documentation for file (source/api_cc/include/custom_op.h
)
#include <vector>
#include <string>
#include <iostream>
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/public/session.h"
#include "tensorflow/core/public/version.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/shape_inference.h"
#include <tensorflow/core/graph/default_device.h>
#include <tensorflow/core/graph/graph_def_builder.h>
using namespace tensorflow;
using CPUDevice = Eigen::ThreadPoolDevice;
using GPUDevice = Eigen::GpuDevice;
// functions used in custom ops
struct DeviceFunctor {
void operator()(
std::string& device,
const CPUDevice& d)
{
device = "CPU";
}
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
void operator()(
std::string& device,
const GPUDevice& d)
{
device = "GPU";
}
#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
};