[doc] add docstr

hpcaitech · Feb 21, 2025 · bec8932 · bec8932
1 parent 4a560d5
commit bec8932
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/applications/ColossalChat/coati/distributed/inference_backend.py b/applications/ColossalChat/coati/distributed/inference_backend.py
@@ -36,7 +36,7 @@ def generate(self, input_ids: torch.Tensor, attention_mask: torch.Tensor, **kwar
                 - attention_mask (torch.Tensor): shape [B, S+N]
                 - action_log_probs (torch.Tensor): shape [B, N]
                 - action_mask (torch.Tensor): shape [B, N]
-                where N is the number of generated tokens.
+                where N is the number of generated tokens. And all tensors should be on CUDA.
         """
 
     def load_state_dict(self, state_dict: Dict[str, torch.Tensor]) -> None: