From 08eba4b78f18e1cd6348bea671e4b12c443a6e81 Mon Sep 17 00:00:00 2001
From: alexcbb <alexchapin@hotmail.fr>
Date: Tue, 18 Jul 2023 16:46:24 +0200
Subject: [PATCH] Fix device

---
 osrt/encoder.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/osrt/encoder.py b/osrt/encoder.py
index 7c09f41..e5292ba 100644
--- a/osrt/encoder.py
+++ b/osrt/encoder.py
@@ -191,7 +191,7 @@ class FeatureMasking(nn.Module):
         im_size = self.resize.apply_image(images[0]).shape[-3:-1]
 
         ### Pre-process images for the image encoder (Resize and Pad)
-        images = torch.stack([self.preprocess(x) for x in images], device=self.mask_generator.device)
+        images = torch.stack([self.preprocess(x) for x in images]).to(self.mask_generator.device)
 
         ### Encode images 
         image_embeddings, embed_no_red = self.mask_generator.image_encoder(images, before_channel_reduc=True) # [B x N, C, H, W] 
-- 
GitLab