Fix sdk error for multi-gpu execution (#2411)

open-mmlab · Sep 8, 2023 · 6e60cae · 6e60cae
1 parent a3c2fae
commit 6e60cae
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 0 deletions.
diff --git a/csrc/mmdeploy/codebase/mmdet/instance_segmentation.cpp b/csrc/mmdeploy/codebase/mmdet/instance_segmentation.cpp
@@ -28,6 +28,7 @@ class ResizeInstanceMask : public ResizeBBox {
   Result<Value> operator()(const Value& prep_res, const Value& infer_res) {
     MMDEPLOY_DEBUG("prep_res: {}\ninfer_res: {}", prep_res, infer_res);
     try {
+      DeviceGuard guard(device_);
       auto dets = infer_res["dets"].get<Tensor>();
       auto labels = infer_res["labels"].get<Tensor>();
       auto masks = infer_res["masks"].get<Tensor>();

diff --git a/csrc/mmdeploy/preprocess/transform/compose.cpp b/csrc/mmdeploy/preprocess/transform/compose.cpp
@@ -55,6 +55,7 @@ class Compose : public Transform {
     if (!hash_code_.empty()) {
       context.set_use_dummy(true);
     }
+    DeviceGuard guard(device_);
     for (size_t i = 0; i < transforms_.size(); ++i) {
       std::optional<profiler::ScopedCounter> child_counter;
       if (scope_) {