pos-egnn

Dockerfile

@@ -4,7 +4,6 @@ WORKDIR /app
 COPY requirements.txt .
 RUN pip install -r --no-cache-dir requirements.txt
 # preload models
-RUN pip install torch_scatter==2.1.2
 RUN python -c '\
 from transformers import BartForConditionalGeneration, AutoTokenizer;\
 AutoTokenizer.from_pretrained("ibm/materials.selfies-ted");\

Dockerfile-conda

@@ -8,7 +8,6 @@ RUN conda create --name fm4m python=3.9.7
 RUN conda activate fm4m
 COPY requirements.txt .
 RUN pip install -r --no-cache-dir requirements.txt
-RUN pip install torch_scatter==2.1.2
 COPY . .
 
 CMD ["python", "app.py"]

models/pos_egnn/load.py

@@ -23,7 +23,8 @@ def smiles_to_atoms(smiles):
     return ase_atoms
 
 class POSEGNN():
-    def __init__(self, device="cpu"):
+    def __init__(self, use_gpu=True):
+        device = "cuda" if use_gpu and torch.cuda.is_available() else "cpu"
         self.device = device
         self.calculator = None
 

requirements.txt

@@ -28,4 +28,4 @@ tqdm>=4.66.4
 pandas==2.2.3
 mordred
 ase==3.24.0
-torch_nl==0.3
+torch_nl==0.3