Fix: correct shape check and torch.save argument order

pooling_generator/token_to_sequence_pooler.py

@@ -19,7 +19,7 @@ def __init__(self, path_token_emb, path_attention_layers, alpha=0.85):
         # remove the CLS and END token embedding
         if len(self.representations_with_cls.shape) == 2:
             self.representations =  self.representations_with_cls[1:-1]
-        elif len(self.representations.shape) == 3:
+        elif len(self.representations_with_cls.shape) == 3:
             self.representations = self.representations_with_cls[:,1:-1, :]
 
         self.attn_all_layers = self._load_torch_data(self.path_attention_layers)
@@ -44,7 +44,7 @@ def cls_pooling(self, save_path=None):
         if self.representations_with_cls is not None:
             cls_token = self.representations_with_cls[0]
             if save_path:
-                torch.save(save_path, cls_token)
+                torch.save(cls_token, save_path)
             return cls_token.squeeze()
         print(f"representations_with_cls was None for sequence {self.uniprot_accession}", flush=True)
         return None  # Handle cases where CLS token is not available or representations are not loaded properly