zerogpu

app.py

@@ -1,6 +1,7 @@
 from pathlib import Path
 from typing import List, Optional, Tuple
 
+import spaces
 import gradio as gr
 import numpy as np
 import torch
@@ -162,7 +163,7 @@ def calculate_surprisal_diff(
     ]
     return diff_tokens2surprisal
 
-
+@spaces.GPU
 def main(input_text: str) -> Tuple[str, str, str]:
     mean_surprisal, char2surprisal = calculate_surprisals_by_character(
         input_text, trained_model, tokenizer