pytorch/caffe2/distributed/file_store_handler_op_gpu.cc
Xiaodong Wang 62d8c8cf0a Manual hipify caffe2/distributed and rocm update (no hcc modules support) (#18088)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/18088

Manually hipify the distributed folder

Reviewed By: bddppq

Differential Revision: D14482702

fbshipit-source-id: cc0abdf525b423ab1f18db8010d21e27c6668d36
2019-03-29 11:07:32 -07:00

21 lines
451 B
C++

#include "caffe2/distributed/file_store_handler_op.h"
#ifndef __HIP_PLATFORM_HCC__
#include <caffe2/core/context_gpu.h>
#else
#include <caffe2/core/hip/context_gpu.h>
#endif
namespace caffe2 {
#ifndef __HIP_PLATFORM_HCC__
REGISTER_CUDA_OPERATOR(
FileStoreHandlerCreate,
FileStoreHandlerCreateOp<CUDAContext>);
#else
REGISTER_HIP_OPERATOR(
FileStoreHandlerCreate,
FileStoreHandlerCreateOp<HIPContext>);
#endif
} // namespace caffe2