Update app.py

app.py

@@ -17,7 +17,6 @@ os.environ['WORLD_SIZE'] = "1"
 os.environ['MASTER_PORT'] = "12193"
 os.environ['MASTER_ADDR'] = "127.0.0.1"
 
-@spaces.GPU
 def load_models():
     model = SuperposedLlama.build(ckpt_dir=weight_path, 
                          tokenizer_path=f'{weight_path}/tokenizer.model', 
@@ -47,6 +46,8 @@ i_weights = params["i_weights"]
 i_length = params["i_length"]
 # Load main model
 model = load_models()
+model.model.to("cuda")
+model.device = "cuda"
 tokenizer = Tokenizer(f'{weight_path}/tokenizer.model')
 # Create ngram models
 ngrams = make_models("ckpts-200k", bigram=True, trigram=True, fourgram=True, fivegram=True, sixgram=True, sevengram=False)
@@ -66,7 +67,6 @@ def decode(tokenizer, encoding):
 
 @spaces.GPU
 def update_options(input, num_tokens):
-    model.to("cuda")
     tokenized_prompts = tokenizer.encode([input], True, False)
     alive_gens, _ = model.sup_generate(prompt_tokens=tokenized_prompts, 
                                             smoothing="geom",