Update tool/comget/generator.py

tool/comget/generator.py

@@ -17,6 +17,10 @@ import sys
 sys.path.append(os.path.join(RDConfig.RDContribDir, 'SA_Score'))
 from rdkit import Chem
 import os
+import os
+import torch
+
+torch.classes.__path__ = [os.path.join(torch.__path__[0], torch.classes.__file__)] 
  
 def get_mol(smiles_or_mol):
     '''
@@ -130,7 +134,7 @@ def generation(value):
         args.data_name = 'ppcenos'
         args.vocab_size = 29  #
         args.block_size = 196 #max_len
-        args.gen_size = 20  
+        args.gen_size = 10  
         args.batch_size = 5
         args.csv_name = 'ppcenos'
         args.props = ['pce']