Tonic commited on
Commit
cc4cdeb
·
1 Parent(s): 7751312

attempt to solve gpu error

Browse files
Files changed (1) hide show
  1. app.py +27 -5
app.py CHANGED
@@ -423,11 +423,33 @@ def make_prediction(symbol: str, timeframe: str = "1d", prediction_days: int = 5
423
  if isinstance(value, torch.Tensor) and hasattr(value, 'to'):
424
  pipe.model.__dict__[name] = value.to(device)
425
 
426
- quantiles, mean = pipe.predict_quantiles(
427
- context=context,
428
- prediction_length=actual_prediction_length,
429
- quantile_levels=[0.1, 0.5, 0.9]
430
- )
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
431
 
432
  if quantiles is None or mean is None:
433
  raise ValueError("Chronos returned empty prediction")
 
423
  if isinstance(value, torch.Tensor) and hasattr(value, 'to'):
424
  pipe.model.__dict__[name] = value.to(device)
425
 
426
+ # Move any additional model components to GPU
427
+ if hasattr(pipe.model, 'config'):
428
+ for key, value in pipe.model.config.__dict__.items():
429
+ if isinstance(value, torch.Tensor) and hasattr(value, 'to'):
430
+ setattr(pipe.model.config, key, value.to(device))
431
+
432
+ # Ensure all model components are in eval mode
433
+ pipe.model.eval()
434
+
435
+ # Move any additional tensors in the pipeline to GPU
436
+ for attr_name in dir(pipe):
437
+ attr = getattr(pipe, attr_name)
438
+ if isinstance(attr, torch.Tensor) and hasattr(attr, 'to'):
439
+ setattr(pipe, attr_name, attr.to(device))
440
+
441
+ # Ensure context is properly shaped and on GPU
442
+ if len(context.shape) == 1:
443
+ context = context.unsqueeze(0)
444
+ context = context.to(device)
445
+
446
+ # Ensure all inputs are on the same device
447
+ with torch.cuda.device(device):
448
+ quantiles, mean = pipe.predict_quantiles(
449
+ context=context,
450
+ prediction_length=actual_prediction_length,
451
+ quantile_levels=[0.1, 0.5, 0.9]
452
+ )
453
 
454
  if quantiles is None or mean is None:
455
  raise ValueError("Chronos returned empty prediction")