Add device to transformer

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, :]