fix

src/main.py

@@ -92,7 +92,7 @@ class StriMap_pHLA:
     def __init__(
         self,
         device: str = 'cuda:0',
-        model_save_path: str = 'model_params/best_model_phla.pt',
+        model_save_path: str = '/data/model_params/best_model_phla.pt',
         pep_dim: int = 256,
         hla_dim: int = 256,
         bilinear_dim: int = 256,
@@ -101,9 +101,9 @@ class StriMap_pHLA:
         gamma: float = 2.0,
         esm2_layer: int = 33,
         batch_size: int = 256,
-        esmfold_cache_dir: str = "esm_cache",
-        cache_dir: str = 'phla_cache',
-        cache_save: bool = True,
+        esmfold_cache_dir: str = "/data/esm_cache",
+        cache_dir: str = '/data/phla_cache',
+        cache_save: bool = False,
         seed: int = 1,
         pos_weights: Optional[float] = None
     ):

src/model.py

@@ -342,7 +342,7 @@ class ESM2Encoder(nn.Module):
     def __init__(self, 
                  device="cuda:0", 
                  layer=33,
-                 cache_dir='cache'):
+                 cache_dir='/data/cache'):
         """
         Initialize an ESM2 encoder.
 
@@ -695,7 +695,7 @@ def batch_embed_to_dicts(
     return emb_dict, coord_dict, failures
 
 class ESMFoldEncoder(nn.Module):
-    def __init__(self, model_name="facebook/esmfold_v1", esm_cache_dir="esm_cache", cache_dir="cache"):
+    def __init__(self, model_name="facebook/esmfold_v1", esm_cache_dir="/data/esm_cache", cache_dir="/data/cache"):
         super(ESMFoldEncoder, self).__init__()
         self.model_name = model_name
         self.esm_cache_dir = esm_cache_dir