Caffe2 - C++ API
A deep learning, cross platform ML framework
load_save_op_gpu.cc
1 
17 #include "caffe2/core/context_gpu.h"
18 #include "caffe2/operators/load_save_op.h"
19 
20 namespace caffe2 {
21 
22 template <>
23 void LoadOp<CUDAContext>::SetCurrentDevice(BlobProto* proto) {
24  if (proto->has_tensor()) {
25  auto* device_detail = proto->mutable_tensor()->mutable_device_detail();
26  device_detail->set_device_type(CUDA);
27  device_detail->set_cuda_gpu_id(CaffeCudaGetDevice());
28  }
29 }
30 
31 REGISTER_CUDA_OPERATOR(Load, LoadOp<CUDAContext>);
32 REGISTER_CUDA_OPERATOR(Save, SaveOp<CUDAContext>);
33 REGISTER_CUDA_OPERATOR(Checkpoint, CheckpointOp<CUDAContext>);
34 } // namespace caffe2
Copyright (c) 2016-present, Facebook, Inc.
int CaffeCudaGetDevice()
Gets the current GPU id.
Definition: common_gpu.cc:125