Skip to content

Commit b2c3941

Browse files
Fix docs for sentence transformer (#10495)
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
1 parent 62d0db3 commit b2c3941

File tree

1 file changed

+13
-0
lines changed

1 file changed

+13
-0
lines changed

torch_geometric/llm/models/sentence_transformer.py

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -112,6 +112,19 @@ def encode(
112112
output_device: Optional[Union[torch.device, str]] = None,
113113
verbose=False,
114114
) -> Tensor:
115+
r"""Main function for users. Converts strings to embeddings.
116+
117+
Args:
118+
text (List[str]): List of strings to embed.
119+
batch_size (int, optional): How many strings to process.
120+
Defaults to processing all at once, but this may lead to
121+
OOM errors. (default: obj:`None`)
122+
output_device (Union[torch.device, str], optional):
123+
By default outputs cpu pytorch tensor, but can choose
124+
to output to specific cuda devices. (default: obj:`None`)
125+
verbose (bool, optional): Controls the verbosity of outputs.
126+
(default: obj:`False`)
127+
"""
115128
is_empty = len(text) == 0
116129
text = ['dummy'] if is_empty else text
117130

0 commit comments

Comments
 (0)