Elgato_dark/chronos-forecasting
1
0
Fork 0
mirror of https://github.com/amazon-science/chronos-forecasting synced 2026-05-23 09:39:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

handle dtype warning

Browse source
This commit is contained in:
Kashif Rasul 2026-01-09 20:44:20 +01:00
parent 16a91d5b58
commit f669e6c5ca
1 changed files with 5 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
src/chronos/base.py
View file

@ -347,9 +347,11 @@ class BaseChronosPipeline(metaclass=PipelineRegistry):
from transformers import AutoConfig
torch_dtype = kwargs.get("torch_dtype", "auto")
if torch_dtype != "auto" and isinstance(torch_dtype, str):
kwargs["torch_dtype"] = cls.dtypes[torch_dtype]
# Handle both torch_dtype (deprecated) and dtype arguments
dtype_value = kwargs.pop("torch_dtype", None) or kwargs.pop("dtype", "auto")
if dtype_value != "auto" and isinstance(dtype_value, str):
dtype_value = cls.dtypes[dtype_value]
kwargs["dtype"] = dtype_value
config = AutoConfig.from_pretrained(pretrained_model_name_or_path, **kwargs)
is_valid_config = hasattr(config, "chronos_pipeline_class") or hasattr(config, "chronos_config")