Update app.py

app.py

@@ -46,7 +46,7 @@ n_token_sample = params["n_token_sample"]
 i_weights = params["i_weights"]
 i_length = params["i_length"]
 # Load main model
-model = load_models().to('cuda')
+model = load_models()
 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,6 +66,7 @@ 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",