revert

lostella · lostella · commit d45c4684156f · 2024-11-06T15:03:13.000+01:00
diff --git a/src/chronos/chronos.py b/src/chronos/chronos.py
@@ -524,7 +524,7 @@ def predict(
             )
             prediction = self.tokenizer.output_transform(
                 samples.to(scale.device), scale
-            )
+            ).to(context_tensor)
 
             predictions.append(prediction)
             remaining -= prediction.shape[-1]
@@ -536,7 +536,7 @@ def predict(
                 [context_tensor, prediction.median(dim=1).values], dim=-1
             )
 
-        return torch.cat(predictions, dim=-1).to(context_tensor)
+        return torch.cat(predictions, dim=-1)
 
     @classmethod
     def from_pretrained(cls, *args, **kwargs):

Original file line number	Diff line number	Diff line change
`@@ -524,7 +524,7 @@ def predict(`
`524`	`524`	`)`
`525`	`525`	`prediction = self.tokenizer.output_transform(`
`526`	`526`	`samples.to(scale.device), scale`
`527`		`- )`
	`527`	`+ ).to(context_tensor)`
`528`	`528`
`529`	`529`	`predictions.append(prediction)`
`530`	`530`	`remaining -= prediction.shape[-1]`
`@@ -536,7 +536,7 @@ def predict(`
`536`	`536`	`[context_tensor, prediction.median(dim=1).values], dim=-1`
`537`	`537`	`)`
`538`	`538`
`539`		`- return torch.cat(predictions, dim=-1).to(context_tensor)`
	`539`	`+ return torch.cat(predictions, dim=-1)`
`540`	`540`
`541`	`541`	`@classmethod`
`542`	`542`	`def from_pretrained(cls, args, *kwargs):`