From d1a1be7889de27097ab54ecf1b25f81d4fcf20ec Mon Sep 17 00:00:00 2001
From: alexcbb <alexchapin@hotmail.fr>
Date: Wed, 19 Jul 2023 10:11:48 +0200
Subject: [PATCH] Add device to transformer
---
osrt/sam/mask_decoder.py | 1 +
1 file changed, 1 insertion(+)
diff --git a/osrt/sam/mask_decoder.py b/osrt/sam/mask_decoder.py
index 9967d16..d0739c8 100644
--- a/osrt/sam/mask_decoder.py
+++ b/osrt/sam/mask_decoder.py
@@ -134,6 +134,7 @@ class MaskDecoder(nn.Module):
b, c, h, w = src.shape
+ self.transformer.to(src.device)
# Run the transformer
hs, src = self.transformer(src, pos_src, tokens)
iou_token_out = hs[:, 0, :]
--
GitLab