smhavens commited on
Commit
5363e91
1 Parent(s): 3922a86

Fix device set up

Browse files
Files changed (1) hide show
  1. app.py +2 -1
app.py CHANGED
@@ -150,7 +150,8 @@ def finetune(train_dataloader):
150
 
151
  def get_model():
152
  model = SentenceTransformer("bert-analogies")
153
- device = torch.device('cuda:0')
 
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