Fix size nn for pocket-conditioned generation

app.py

@@ -465,7 +465,7 @@ def generate_with_pocket(in_fragments, in_protein, n_steps, n_atoms, num_samples
 
     if n_atoms == 0:
         def sample_fn(_data):
-            out, _ = size_nn.forward(_data, return_loss=False)
+            out, _ = size_nn.forward(_data, return_loss=False, with_pocket=True)
             probabilities = torch.softmax(out, dim=1)
             distribution = torch.distributions.Categorical(probs=probabilities)
             samples = distribution.sample()
@@ -536,7 +536,7 @@ with demo:
                     label="Number of Denoising Steps", step=10
                 )
                 n_atoms = gr.Slider(
-                    minimum=0, maximum=20, value=5,
+                    minimum=0, maximum=20,
                     label="Linker Size: DiffLinker will predict it if set to 0",
                     step=1
                 )

src/linker_size_lightning.py

@@ -79,10 +79,10 @@ class SizeClassifier(pl.LightningModule):
     def test_dataloader(self):
         return get_dataloader(self.test_dataset, self.batch_size, collate_fn=collate_with_fragment_edges)
 
-    def forward(self, data, return_loss=True):
+    def forward(self, data, return_loss=True, with_pocket=False):
         h = data['one_hot']
         x = data['positions']
-        fragment_mask = data['fragment_mask']
+        fragment_mask = data['fragment_only_mask'] if with_pocket else data['fragment_mask']
         linker_mask = data['linker_mask']
         edge_mask = data['edge_mask']
         edges = data['edges']