SalesforceAIResearch / uni2ts

[ICML2024] Unified Training of Universal Time Series Forecasting Transformers
Apache License 2.0
611 stars 49 forks source link

GPU, single-node multi-GPU, multi-node multi-GPU support for forecast generation #62

Open ryuta-yoshimatsu opened 1 month ago

ryuta-yoshimatsu commented 1 month ago

Is your feature request related to a problem? Please describe. It is not clear from the documentation and the sample code, if the forecast generation can be performed on a GPU, multiple GPUs, or multiple GPUs in multiple nodes. If this is the case, please add some descriptions on how to achieve this.

Describe the solution you'd like Maybe a simple device_map="auto" argument support in the from_pretrained method that automatically loads the model and the subsequent data onto a GPU/s.

Additional context I respect the great work you have done! Especially around enabling multivariate modelling.