Update README.md

README.md

@@ -26,6 +26,11 @@ model = AutoModelForCausalLM.from_pretrained(
 batch_size, lookback_length, channels = 1, 2880, 7
 time_series = torch.randn(batch_size, lookback_length, channels)
 
+# Load the model and data to the same device
+device = torch.cuda.current_device() if torch.cuda.is_available() else 'cpu'
+model = model.to(device)
+time_series = time_series.to(device)
+
 # Generate forecasts
 forecast_length = 96
 predictions = model.generate(time_series, max_new_tokens=forecast_length)