Update tool/comget/generator.py

tool/comget/generator.py

@@ -17,8 +17,7 @@ import sys
 sys.path.append(os.path.join(RDConfig.RDContribDir, 'SA_Score'))
 from rdkit import Chem
 import os
-os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
-
+ 
 def get_mol(smiles_or_mol):
     '''
     Loads SMILES/molecule into RDKit's object
@@ -206,9 +205,9 @@ def generation(value):
                         x = torch.tensor([stoi[s] for s in regex.findall(context)], dtype=torch.long)[None,...].repeat(args.batch_size, 1).to('cuda')
                         p = None
                         if len(args.props) == 1:
-                                p = torch.tensor([c]).repeat(args.batch_size, 1).to('cuda')   # for single condition
+                                p = torch.tensor([c]).repeat(args.batch_size, 1).to('cpu')   # for single condition
                         else:
-                                p = torch.tensor([c]).repeat(args.batch_size, 1).unsqueeze(1).to('cuda')    # for multiple conditions
+                                p = torch.tensor([c]).repeat(args.batch_size, 1).unsqueeze(1).to('cpu')    # for multiple conditions
                         sca = None
                         y = sample(model, x, 300, temperature= 1.0, sample=True, top_k = 10, prop = p, scaffold = sca)   
                         for gen_mol in y: