Almaatla commited on
Commit
536f938
1 Parent(s): 288dd42

add cuda and gpu space

Browse files
Files changed (1) hide show
  1. classification.py +2 -2
classification.py CHANGED
@@ -133,9 +133,9 @@ def load_data(file_obj):
133
  # Assuming file_obj is a file-like object uploaded via Gradio, use `pd.read_excel` directly on it
134
  return pd.read_excel(file_obj)
135
 
136
-
137
  def initialize_models():
138
- model_ST = SentenceTransformer("all-mpnet-base-v2")
139
  return model_ST
140
 
141
 
 
133
  # Assuming file_obj is a file-like object uploaded via Gradio, use `pd.read_excel` directly on it
134
  return pd.read_excel(file_obj)
135
 
136
+ @spaces.GPU
137
  def initialize_models():
138
+ model_ST = SentenceTransformer("all-mpnet-base-v2",device = "cuda")
139
  return model_ST
140
 
141