Fix an issue with working with gpu

The variable ab needs to be registered to buffer to avoid an error with incompatible devices.
ChaoningZhang · Jul 2, 2023 · 48f418b · 48f418b
1 parent 42e87b2
commit 48f418b
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/mobile_sam/modeling/tiny_vit_sam.py b/mobile_sam/modeling/tiny_vit_sam.py
@@ -253,7 +253,9 @@ def train(self, mode=True):
         if mode and hasattr(self, 'ab'):
             del self.ab
         else:
-            self.ab = self.attention_biases[:, self.attention_bias_idxs]
+            self.register_buffer('ab',
+                                 self.attention_biases[:, self.attention_bias_idxs],
+                                 persistent=False)
 
     def forward(self, x):  # x (B,N,C)
         B, N, _ = x.shape