Update README.md

README.md

@@ -16,11 +16,11 @@ You can find more details about the model on [GitHub page](https://github.com/an
 ## 🚀 Quick Start
 ```python
 import torch
-from transformers import AutoModelForCausalLM, AutoConfig
+from transformers import AutoModel
 
 # Load pre-trained model (when available)
-model = AutoModelForCausalLM.from_pretrained(
-    'ant-intl/Falcon-TST_Large', 
+model = AutoModel.from_pretrained(
+    'ant-intl/Falcon-TST_Large',
     trust_remote_code=True
 )
 
@@ -32,8 +32,9 @@ time_series = torch.randn(batch_size, lookback_length, channels)
 device = torch.cuda.current_device() if torch.cuda.is_available() else 'cpu'
 model = model.to(device)
 time_series = time_series.to(device)
+print(f'device: {device}')
 
 # Generate forecasts
 forecast_length = 96
-predictions = model.generate(time_series, max_new_tokens=forecast_length)
+predictions = model.predict(time_series, forecast_horizon=forecast_length)
 ```