Skip to content

Commit 357752f

Browse files
committed
esm2 reader: save reader to default global data dir
1 parent ad24fa7 commit 357752f

1 file changed

Lines changed: 2 additions & 2 deletions

File tree

chebai/preprocessing/reader.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -509,7 +509,7 @@ class ESM2EmbeddingReader(DataReader):
509509

510510
def __init__(
511511
self,
512-
save_model_dir: str,
512+
save_model_dir: str = os.path.join("data", "esm2_reader"),
513513
model_name: str = "esm2_t36_3B_UR50D",
514514
device: Optional[torch.device] = None,
515515
truncation_length: int = 1022,
@@ -617,7 +617,7 @@ def load_hub_workaround(self, url) -> torch.Tensor:
617617
return data
618618

619619
@staticmethod
620-
def name() -> None:
620+
def name() -> str:
621621
"""
622622
Returns the name of the data reader. This method identifies the specific type of data reader.
623623

0 commit comments

Comments
 (0)