Fix norm3

osrt/layers.py

@@ -261,3 +261,7 @@ class SlotAttention(nn.Module):
     def change_slots_number(self, num_slots): 
         self.num_slots = num_slots
         self.initial_slots = nn.Parameter(torch.randn(num_slots, self.slot_dim))
+
+
+class TransformerSlotAttention(nn.Module):
+    pass
\ No newline at end of file

osrt/sam/transformer.py

@@ -174,8 +174,8 @@ class TwoWayAttentionBlock(nn.Module):
         self.mlp = self.mlp.to(queries.device)
         mlp_out = self.mlp(queries)
         queries = queries + mlp_out
-        queries = self.norm3(queries)
         self.norm3 = self.norm3.to(queries.device)
+        queries = self.norm3(queries)
 
         # Cross attention block, image embedding attending to tokens
         q = queries + query_pe