1 #include "caffe2/core/context_gpu.h" 2 #include "caffe2/operators/load_save_op.h" 7 void LoadOp<CUDAContext>::SetCurrentDevice(BlobProto* proto) {
8 if (proto->has_tensor()) {
9 auto* device_detail = proto->mutable_tensor()->mutable_device_detail();
10 device_detail->set_device_type(CUDA);
15 REGISTER_CUDA_OPERATOR(Load, LoadOp<CUDAContext>);
16 REGISTER_CUDA_OPERATOR(Save, SaveOp<CUDAContext>);
17 REGISTER_CUDA_OPERATOR(Checkpoint, CheckpointOp<CUDAContext>);
A global dictionary that holds information about what Caffe2 modules have been loaded in the current ...
int CaffeCudaGetDevice()
Gets the current GPU id.