Spaces:
Sleeping
Sleeping
smhavens
commited on
Commit
•
5363e91
1
Parent(s):
3922a86
Fix device set up
Browse files
app.py
CHANGED
@@ -150,7 +150,8 @@ def finetune(train_dataloader):
|
|
150 |
|
151 |
def get_model():
|
152 |
model = SentenceTransformer("bert-analogies")
|
153 |
-
|
|
|
154 |
model = model.to(device)
|
155 |
return model
|
156 |
|
|
|
150 |
|
151 |
def get_model():
|
152 |
model = SentenceTransformer("bert-analogies")
|
153 |
+
gpu_available = torch.cuda.is_available()
|
154 |
+
device = torch.device("cuda" if gpu_available else "cpu")
|
155 |
model = model.to(device)
|
156 |
return model
|
157 |
|