update

utils.py

@@ -279,10 +279,13 @@ class MolecularGenerationModel():
             input_length = batch['input_ids'].shape[1]
             steps = 1024 - input_length
 
+            print(self.model.device, "model_device")
             with torch.set_grad_enabled(False):
                 early_stop_flags = torch.zeros(num_generations, dtype=torch.bool).to(self.model.device)
                 for k in range(steps):
+                    print("batch", batch)
                     logits = self.model(**batch)['logits']
+                    print("logits", logits)
                     logits = logits[:, -1, :] / temperature
                     probs = F.softmax(logits, dim=-1)
                     ix = torch.multinomial(probs, num_samples=num_generations)