123

app.py

@@ -3,6 +3,7 @@ import gradio as gr
 from rdkit import Chem
 from rdkit.Chem import Draw
 import os
+import spaces
 import torch
 try:
     print(f"Is CUDA available: {torch.cuda.is_available()}")
@@ -10,7 +11,8 @@ try:
     print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
     # Tesla T4
 except:
-    pass
+    print("No GPU found!")
+
 
 def remove_atom_mapping_and_isotopes(smiles):
     mol = Chem.MolFromSmiles(smiles)

infer.sh

@@ -7,4 +7,5 @@ python inference.py \
 --batch_size 16384 \
 --batch_type tokens \
 --max_length 500 \
---seed 0
+--seed 0 \
+--gpu 0

onmt/bin/translate.py

@@ -9,6 +9,7 @@ from onmt.constants import CorpusTask
 import onmt.opts as opts
 from onmt.utils.parse import ArgumentParser
 from onmt.utils.misc import use_gpu, set_random_seed
+import spaces
 
 
 def translate(opt):
@@ -49,7 +50,7 @@ def _get_parser():
     opts.translate_opts(parser, dynamic=True)
     return parser
 
-
+@spaces.GPU
 def main():
     parser = _get_parser()
     opt = parser.parse_args()