Elgato_dark/chronos-forecasting
1
0
Fork 0
mirror of https://github.com/amazon-science/chronos-forecasting synced 2026-05-24 01:58:27 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Chronos-2: Only pin_memory when device type is cuda (#431)

Browse source
This commit is contained in:
Abdul Fatir 2025-12-10 14:59:13 +01:00 committed by GitHub
parent 086e660023
commit eb5b61234a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 9 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
src/chronos/chronos2/pipeline.py
View file

@ -608,7 +608,9 @@ class Chronos2Pipeline(BaseChronosPipeline):
output_patch_size=self.model_output_patch_size,
mode=DatasetMode.TEST,
)
test_loader = DataLoader(test_dataset, batch_size=None, pin_memory=True, shuffle=False, drop_last=False)
test_loader = DataLoader(
test_dataset, batch_size=None, pin_memory=self.model.device.type == "cuda", shuffle=False, drop_last=False
)
all_predictions: list[torch.Tensor] = []
for batch in test_loader:
@ -1122,7 +1124,12 @@ class Chronos2Pipeline(BaseChronosPipeline):
mode=DatasetMode.TEST,
)
test_loader = DataLoader(
test_dataset, batch_size=None, num_workers=1, pin_memory=True, shuffle=False, drop_last=False
test_dataset,
batch_size=None,
num_workers=0,
pin_memory=self.model.device.type == "cuda",
shuffle=False,
drop_last=False,
)
all_embeds: list[torch.Tensor] = []
all_loc_scales: list[tuple[torch.Tensor, torch.Tensor]] = []