try not pooling all the sos tokens and instead using the last one

Author: lucidrains

meshgpt_pytorch/meshgpt_pytorch.py

@@ -1475,10 +1475,10 @@ def forward_on_codes(
 
         # if calling without kv cache, pool the sos tokens, if greater than 1 sos token
 
-        if not exists(cache) and self.num_sos_tokens > 1:
+        if not exists(cache):
             sos_tokens, attended_face_codes = unpack(attended_face_codes, packed_sos_shape, 'b * d')
-            pooled_sos_token = reduce(sos_tokens, 'b n d -> b 1 d', 'mean')
-            attended_face_codes = torch.cat((pooled_sos_token, attended_face_codes), dim = 1)
+            last_sos_token = sos_tokens[:, -1:]
+            attended_face_codes = torch.cat((last_sos_token, attended_face_codes), dim = 1)
 
         # maybe project from coarse to fine dimension for hierarchical transformers
 

meshgpt_pytorch/version.py

@@ -1 +1 @@
-__version__ = '1.2.5'
+__version__ = '1.2.6'