Function prepare_batch

tests/models/patchtst/test_modeling_patchtst.py:304–308

(repo_id="hf-internal-testing/etth1-hourly-batch", file="train-batch.pt")

def prepare_batch(repo_id="hf-internal-testing/etth1-hourly-batch", file="train-batch.pt"):
    file = hf_hub_download(repo_id=repo_id, filename=file, repo_type="dataset")
    check_torch_load_is_safe()
    batch = torch.load(file, map_location=torch_device, weights_only=True)
    return batch


# Note: Pretrained model is not yet downloadable.