Address comments

amazon-science · Mar 24, 2024 · 6ffc3ba · 6ffc3ba
1 parent 763540a
commit 6ffc3ba
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/src/chronos/chronos.py b/src/chronos/chronos.py
@@ -209,7 +209,6 @@ def __init__(self, config: ChronosConfig, model: PreTrainedModel) -> None:
     def device(self):
         return self.model.device
 
-    @torch.no_grad()
     def encode(
         self,
         input_ids: torch.Tensor,
@@ -349,10 +348,12 @@ def _prepare_and_validate_context(
 
         return context
 
+    @torch.no_grad()
     def embed(
         self, context: Union[torch.Tensor, List[torch.Tensor]]
     ) -> Tuple[torch.Tensor, torch.Tensor]:
-        """Get encoder embeddings for the given time series.
+        """
+        Get encoder embeddings for the given time series.
 
         Parameters
         ----------