kz209 commited on
Commit
8acdc64
·
1 Parent(s): b9da634

no gpu at initialization

Browse files
Files changed (1) hide show
  1. pages/summarization_playground.py +4 -3
pages/summarization_playground.py CHANGED
@@ -36,9 +36,10 @@ def generate_answer(sources, model, model_name, prompt):
36
  global __model_on_gpu__
37
 
38
  if __model_on_gpu__ != model_name:
39
- model[__model_on_gpu__].cpu()
40
- gc.collect()
41
- torch.cuda.empty_cache()
 
42
 
43
  model[model_name].cuda()
44
  __model_on_gpu__ = model_name
 
36
  global __model_on_gpu__
37
 
38
  if __model_on_gpu__ != model_name:
39
+ if not __model_on_gpu__:
40
+ model[__model_on_gpu__].cpu()
41
+ gc.collect()
42
+ torch.cuda.empty_cache()
43
 
44
  model[model_name].cuda()
45
  __model_on_gpu__ = model_name