From 86ac08b5d415128eadb550e9e994549a0acc7841 Mon Sep 17 00:00:00 2001
From: alexcbb <alexchapin@hotmail.fr>
Date: Wed, 19 Jul 2023 13:13:46 +0200
Subject: [PATCH] Fix out proj
---
osrt/sam/transformer.py | 1 +
1 file changed, 1 insertion(+)
diff --git a/osrt/sam/transformer.py b/osrt/sam/transformer.py
index 4d88344..a39d564 100644
--- a/osrt/sam/transformer.py
+++ b/osrt/sam/transformer.py
@@ -238,6 +238,7 @@ class Attention(nn.Module):
# Get output
out = attn @ v
out = self._recombine_heads(out)
+ self.out_proj = self.out_proj.to(out.device)
out = self.out_proj(out)
return out
--
GitLab