gpu inference

app.py

@@ -16,6 +16,9 @@ DEBIASING_KEYWORDS = [
     "(rude) ", "(sexually explicit) ", "(hateful) ", "(aggressive) ", "(racist) ", "(threat) ", "(violent) ", "(sexist) "
 ]
 
+if torch.cuda.is_available():
+    use_cuda = True
+
 def debias(prompt, model,use_prefix, max_length=50, num_beam=3):
     """
     Debiasing inference function.
@@ -24,7 +27,7 @@ def debias(prompt, model,use_prefix, max_length=50, num_beam=3):
     :param max_length: The maximum length of the output sentence.
     :return: The debiased output sentence.
     """
-    wrapper = GPT2Wrapper(model_name=str(model), use_cuda=False)
+    wrapper = GPT2Wrapper(model_name=str(model), use_cuda=use_cuda)
     if use_prefix == 'Prefixes':
         debiasing_prefixes = DEBIASING_PREFIXES
     else: