diff options
Diffstat (limited to 'caffe2/utils')
| -rw-r--r-- | caffe2/utils/proto_utils.cc | 4 | ||||
| -rw-r--r-- | caffe2/utils/proto_utils_test.cc | 4 |
2 files changed, 4 insertions, 4 deletions
diff --git a/caffe2/utils/proto_utils.cc b/caffe2/utils/proto_utils.cc index dc8e088eba..dd80282238 100644 --- a/caffe2/utils/proto_utils.cc +++ b/caffe2/utils/proto_utils.cc @@ -30,7 +30,7 @@ C10_EXPORT int DeviceId(const DeviceOption& option) { case PROTO_CPU: return option.numa_node_id(); case PROTO_CUDA: - return option.cuda_gpu_id(); + return option.device_id(); case PROTO_MKLDNN: return option.numa_node_id(); case PROTO_HIP: @@ -43,7 +43,7 @@ C10_EXPORT int DeviceId(const DeviceOption& option) { C10_EXPORT bool IsSameDevice(const DeviceOption& lhs, const DeviceOption& rhs) { return ( lhs.device_type() == rhs.device_type() && - lhs.cuda_gpu_id() == rhs.cuda_gpu_id() && + lhs.device_id() == rhs.device_id() && lhs.hip_gpu_id() == rhs.hip_gpu_id() && lhs.node_name() == rhs.node_name() && lhs.numa_node_id() == rhs.numa_node_id()); diff --git a/caffe2/utils/proto_utils_test.cc b/caffe2/utils/proto_utils_test.cc index c9f37f4c98..5d8fb86b34 100644 --- a/caffe2/utils/proto_utils_test.cc +++ b/caffe2/utils/proto_utils_test.cc @@ -11,9 +11,9 @@ TEST(ProtoUtilsTest, IsSameDevice) { EXPECT_FALSE(IsSameDevice(a, b)); b.set_node_name("my_node"); EXPECT_TRUE(IsSameDevice(a, b)); - b.set_cuda_gpu_id(2); + b.set_device_id(2); EXPECT_FALSE(IsSameDevice(a, b)); - a.set_cuda_gpu_id(2); + a.set_device_id(2); EXPECT_TRUE(IsSameDevice(a, b)); a.set_device_type(DeviceTypeProto::PROTO_CUDA); b.set_device_type(DeviceTypeProto::PROTO_CPU); |