add cuda and gpu space
Browse files- 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 |
|