Skip to content

Commit 6e60cae

Browse files
authored
Fix sdk error for multi-gpu execution (#2411)
1 parent a3c2fae commit 6e60cae

File tree

2 files changed

+2
-0
lines changed

2 files changed

+2
-0
lines changed

csrc/mmdeploy/codebase/mmdet/instance_segmentation.cpp

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@ class ResizeInstanceMask : public ResizeBBox {
2828
Result<Value> operator()(const Value& prep_res, const Value& infer_res) {
2929
MMDEPLOY_DEBUG("prep_res: {}\ninfer_res: {}", prep_res, infer_res);
3030
try {
31+
DeviceGuard guard(device_);
3132
auto dets = infer_res["dets"].get<Tensor>();
3233
auto labels = infer_res["labels"].get<Tensor>();
3334
auto masks = infer_res["masks"].get<Tensor>();

csrc/mmdeploy/preprocess/transform/compose.cpp

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,7 @@ class Compose : public Transform {
5555
if (!hash_code_.empty()) {
5656
context.set_use_dummy(true);
5757
}
58+
DeviceGuard guard(device_);
5859
for (size_t i = 0; i < transforms_.size(); ++i) {
5960
std::optional<profiler::ScopedCounter> child_counter;
6061
if (scope_) {

0 commit comments

Comments
 (0)