Fix: correct shape check and torch.save argument order#1

pawelpratyush wants to merge 1 commit into

Helix-Research-Lab:mainHelix-Research-Lab/Pool_PaRTI:mainfrom

pawelpratyush:fix/pooler-shape-and-save-bugspawelpratyush/Pool_PaRTI:fix/pooler-shape-and-save-bugsCopy head branch name to clipboard