Added cuda

app.py

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