From 1b23903e3dd275d0327c86be148b151d3b0d9d4a Mon Sep 17 00:00:00 2001
From: pawelparker <pawel.pavi@gmail.com>
Date: Mon, 27 Apr 2026 20:01:56 -0400
Subject: [PATCH] Fix: correct shape check and torch.save argument order

---
 .DS_Store                                     | Bin 0 -> 8196 bytes
 pooling_generator/token_to_sequence_pooler.py |   4 ++--
 protein_tasks/.DS_Store                       | Bin 0 -> 8196 bytes
 3 files changed, 2 insertions(+), 2 deletions(-)
 create mode 100644 .DS_Store
 create mode 100644 protein_tasks/.DS_Store

diff --git a/.DS_Store b/.DS_Store
new file mode 100644
index 0000000000000000000000000000000000000000..a88737f770f4f8ee0c6f688ff7c0f4a8cecd8db8
GIT binary patch
literal 8196
zcmeHMF-s#s6n>*A0*4Sm?sDD!01?s%cB|w-tPbpKa*5G3m}J8yS6sK^m3H?B?1a0<
zKVU6fW#b>P7c8%m``*kLXIzcz3ZnQX%zT@9^Jd=n%_N(*AtJT%y<H=!5>XkIvGM`U
zJ%#ODGiA#5u(N_cb)(w)X51TOMoNb$APR^AqJStM3cMBtaL?vae&D`utkQ}CqQJXU
zfWIFCRL0O^YEXZ5pwTG+u!L@1a9!aGU|=z{m>PryMHyG1aaA_OP{tkofgKlGObr@$
zQa1Tec4cKV6s1?k{DB@$Dl{mqC?E=C72w>xL<<yAhelrhzCFM1+In%l-s;9J^p#JJ
zHm0j*r^~**J<&ecah?rChj67i+N6ZGs7WKN-bV}#+Qz#f-Z%Aihq|y2y|`%M{3*U}
z+}8(t|2~<xbFbZtZ+R8*_jpstyk}`ZF<#v<>j*6SsEO7hro?BQ9_90LxqkgU=9*(Y
zm*vwueSXKxZwCAl@Ed^<?`uQv=a6>q?Br^HY4T{7&&SL0nVT@b%OjtM7ysH7$Nw4e
z#)y{pq=B|aEwpXCZcgbV`?YahY!;XApTGMzXSF~)t`5(e$<qfLn_<?0K75`Ny~}-k
zrnjJ46WTR#!hhh?+_|&wi_x!Qv$*<Poc+9CpIq^{`aEwYPoKp@)<=pc@D>!9HbGW7
z|Nrvv^Z&O%MPd{MM1g_|s8W5sz5-{PTW6-3LQjDD4V4RzOATrlG(3d+bvh0l|6z#h
bz}iBKsX<s!Ui?9T{yif1e=WW=?^c2T8xLH~

literal 0
HcmV?d00001

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 0000000000000000000000000000000000000000..ace41c15f0f67c56ce0c03ee7da31e1084ddfd89
GIT binary patch
literal 8196
zcmeI1&ubGw6vyAhf&o#gAT4<4g8l>XFWA-;JqQMJmbUpp8`xw+ng-g7Yc2Rkc=12*
z=1uVCQSt1-Ud4k)kAB~q>1M_lNH#@+z9}<rcXr;}c^}@)<ZX$_Y#p^KMC(LUMrFBr
z4b2}4Kj%78O0I{M4fKf;+M`4K?a*Liw00-~C7=Y9fD%vw%Rm5UHkb00bKgGIRtYG9
z|B?Xj4-qQM*vZUF**efj3IMr+ZYj7%9w0iAld+STm7)rpHhVBm)wmSHSam#)csOM2
zWM-vRCu7ygxHlV@p%{5P^dcEfCbm*rC7=W*32^Ryi?-<ja`tABzqe_KEWSf+8X>Ds
zlKo1(nRfeP!hE?pt6n#)RvSs$!1!|I&GVmYpFV8PuG=cC+duH<`xURG4d-VcezmCu
z-<t5TLn-y({YbuBJjVS#9iV>`b%=4ji!mXA)7Ohz^x3>T&)F<E&u%)Vc6nAqp6}s2
z;)dy*&@pupAxyFfj(d*k2=_xe4ydiZ`#h_j&*tpu={Rm?c&ATbUfsaYn-cDKh$AMI
zXy{0+{C>~994mv|d&E@W`Sn^Io)@&_vv~F~XHnSNI6fA)7Vx?ur|)g5V?TL9HAF)K
zDkIc{?&FTv-Sl0h^nDGC=Fs-JJo=WepSdiISroGn_+R<~2k<fVkdCN<b-;cnxaRMo
z1EG<k7rQCe#tEd`eS%22`|4p48$OHP+ZPw>W$wB7YiAKV0hfC<=Avb}2wbt#S?BzJ
zd;0hP<(jn4MhPf^OC?}ds`cs?OuF29!Z_#J5$Y}~7ao^cDJp0r9fy>395VlhA+93~
jWu4f`%u4j2{Phn3!S4_MwA1sSGa;Xdp8pb}|77zE)n3wM

literal 0
HcmV?d00001