updates

app.py

@@ -1,4 +1,5 @@
 import gradio as gr
+import numpy as np
 import os
 import torch
 import subprocess
@@ -109,9 +110,13 @@ def generate(input_file):
     except Exception as e:
         return f'Could not read the molecule: {e}'
 
+    if molecule.GetNumAtoms() > 50:
+        return f'Too large molecule: upper limit is 50 heavy atoms'
+
     positions, one_hot, charges = parse_molecule(molecule, is_geom=True)
-    positions = torch.tensor(positions, dtype=const.TORCH_FLOAT, device=device)
-    one_hot = torch.tensor(one_hot, dtype=const.TORCH_FLOAT, device=device)
+    anchors = np.zeros_like(charges)
+    fragment_mask = np.ones_like(charges)
+    linker_mask = np.zeros_like(charges)
     print('Read and parsed molecule')
 
     dataset = [{
@@ -120,9 +125,9 @@ def generate(input_file):
         'positions': torch.tensor(positions, dtype=const.TORCH_FLOAT, device=device),
         'one_hot': torch.tensor(one_hot, dtype=const.TORCH_FLOAT, device=device),
         'charges': torch.tensor(charges, dtype=const.TORCH_FLOAT, device=device),
-        'anchors': torch.zeros_like(charges, dtype=const.TORCH_FLOAT, device=device),
-        'fragment_mask': torch.ones_like(charges, dtype=const.TORCH_FLOAT, device=device),
-        'linker_mask': torch.zeros_like(charges, dtype=const.TORCH_FLOAT, device=device),
+        'anchors': torch.tensor(anchors, dtype=const.TORCH_FLOAT, device=device),
+        'fragment_mask': torch.tensor(fragment_mask, dtype=const.TORCH_FLOAT, device=device),
+        'linker_mask': torch.tensor(linker_mask, dtype=const.TORCH_FLOAT, device=device),
         'num_atoms': len(positions),
     }]
     dataloader = get_dataloader(dataset, batch_size=1, collate_fn=collate_with_fragment_edges)