Update pipe.py

pipe.py

@@ -72,7 +72,7 @@ class SmilesDiffusionPipe(Pipeline):
         return torch.multinomial(probs, 1)
 
 
-    def unmask_partial_smiles(self, input_ids, model, voc, steps=None):
+    def unmask_partial_smiles(self, input_ids, model, voc, steps=None, k=None):
         model.eval()
         with torch.no_grad():
             sequences = torch.tensor(input_ids, dtype=torch.long, device=self.device).unsqueeze(0)