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