Update models/peptide_classifiers.py

models/peptide_classifiers.py

@@ -530,9 +530,9 @@ class CrossAttnUnpooled(nn.Module):
         h = self.shared(z)
         return self.reg(h).squeeze(-1), self.cls(h)
 
-def load_affinity_predictor(checkpoint_path, device):
+def load_affinity_predictor(device):
     """Load trained model from checkpoint."""
-    checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
+    checkpoint = torch.load('./classifier_ckpt/wt_affinity.pt', map_location=device, weights_only=False)
 
     model = CrossAttnUnpooled()