diff --git a/.DS_Store b/.DS_Store
new file mode 100644
index 0000000..a88737f
Binary files /dev/null and b/.DS_Store differ
diff --git a/pooling_generator/token_to_sequence_pooler.py b/pooling_generator/token_to_sequence_pooler.py
index 66afb83..944f6f3 100644
--- a/pooling_generator/token_to_sequence_pooler.py
+++ b/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
diff --git a/protein_tasks/.DS_Store b/protein_tasks/.DS_Store
new file mode 100644
index 0000000..ace41c1
Binary files /dev/null and b/protein_tasks/.DS_Store differ