kz209 commited on
Commit
57ace24
β€’
1 Parent(s): c689aee

delete to change models

Browse files
Files changed (1) hide show
  1. pages/summarization_playground.py +3 -3
pages/summarization_playground.py CHANGED
@@ -11,7 +11,7 @@ import torch
11
  load_dotenv()
12
 
13
  __model_on_gpu__ = Model.__model_list__[0]
14
- model = {model_name: Model(model_name) for model_name in Model.__model_list__}
15
 
16
  random_label = 'πŸ”€ Random dialogue from dataset'
17
  examples = {
@@ -37,11 +37,11 @@ def generate_answer(sources, model, model_name, prompt):
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
46
 
47
  answer = model[model_name].gen(content)
 
11
  load_dotenv()
12
 
13
  __model_on_gpu__ = Model.__model_list__[0]
14
+ model = {model_name: None for model_name in Model.__model_list__}
15
 
16
  random_label = 'πŸ”€ Random dialogue from dataset'
17
  examples = {
 
37
 
38
  if __model_on_gpu__ != model_name:
39
  if not __model_on_gpu__:
40
+ del model[__model_on_gpu__]
41
  gc.collect()
42
  torch.cuda.empty_cache()
43
 
44
+ model[model_name] = Model(model_name)
45
  __model_on_gpu__ = model_name
46
 
47
  answer = model[model_name].gen(content)