fix data type

Project-MONAI · Jul 12, 2024 · 80a3ace · 80a3ace
1 parent 6448762
commit 80a3ace
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/scripts/utils/trans_utils.py b/scripts/utils/trans_utils.py
@@ -349,7 +349,7 @@ def __call__(
                     pred += 0.5  # inplace mapping to avoid cloning pred
                     for i in range(1, object_num + 1):
                         frac = i + 0.5
-                        pred[pred == frac] = data["label_prompt"][i - 1].to(pred.dtype)
+                        pred[pred == frac] = torch.tensor(data["label_prompt"][i - 1]).to(pred.dtype)
                     pred[pred == 0.5] = 0.0
                 data[keys] = pred
         return data