Nicolai Berk commited on
Commit
8ca6390
·
1 Parent(s): a0bc956

Add GPU decorator for spaces

Browse files
Files changed (1) hide show
  1. app.py +9 -1
app.py CHANGED
@@ -5,8 +5,16 @@ import faiss
5
  import torch
6
  import numpy as np
7
  from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
8
-
9
  import os
 
 
 
 
 
 
 
 
 
10
 
11
  # Login automatically if HF_TOKEN is present
12
  hf_token = os.getenv("HF_TOKEN")
 
5
  import torch
6
  import numpy as np
7
  from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 
8
  import os
9
+ import spaces
10
+
11
+ @spaces.GPU
12
+ def claim_gpu():
13
+ # Dummy function to make Spaces detect GPU usage
14
+ pass
15
+
16
+ claim_gpu()
17
+
18
 
19
  # Login automatically if HF_TOKEN is present
20
  hf_token = os.getenv("HF_TOKEN")