Trương Gia Bảo commited on
Commit
04d1b54
·
1 Parent(s): c274248

Update GPU app

Browse files
Files changed (4) hide show
  1. app.py +1 -0
  2. requirements-cpu.txt +12 -0
  3. requirements-gpu.txt +13 -0
  4. requirements.txt +2 -1
app.py CHANGED
@@ -189,6 +189,7 @@ def load_model():
189
  model = JTPropVAE(vocab, 450, 56, 20, 3)
190
  if torch.cuda.is_available():
191
  model.load_state_dict(torch.load(model_path))
 
192
  else:
193
  model.load_state_dict(torch.load(model_path,map_location=torch.device('cpu')))
194
  return model
 
189
  model = JTPropVAE(vocab, 450, 56, 20, 3)
190
  if torch.cuda.is_available():
191
  model.load_state_dict(torch.load(model_path))
192
+ model.to('cuda')
193
  else:
194
  model.load_state_dict(torch.load(model_path,map_location=torch.device('cpu')))
195
  return model
requirements-cpu.txt ADDED
@@ -0,0 +1,12 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ torch
2
+ rdkit
3
+ numpy
4
+ tqdm
5
+ networkx
6
+ scipy
7
+ molbloom
8
+ stqdm
9
+ pandas
10
+ streamlit_ext
11
+ streamlit_lottie
12
+ requests
requirements-gpu.txt ADDED
@@ -0,0 +1,13 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ --extra-index-url https://download.pytorch.org/whl/cu113
2
+ torch
3
+ rdkit
4
+ numpy
5
+ tqdm
6
+ networkx
7
+ scipy
8
+ molbloom
9
+ stqdm
10
+ pandas
11
+ streamlit_ext
12
+ streamlit_lottie
13
+ requests
requirements.txt CHANGED
@@ -1,6 +1,7 @@
 
 
1
  rdkit
2
  numpy
3
- torch
4
  tqdm
5
  networkx
6
  scipy
 
1
+ --extra-index-url https://download.pytorch.org/whl/cu113
2
+ torch
3
  rdkit
4
  numpy
 
5
  tqdm
6
  networkx
7
  scipy