Update inference.py

inference.py

@@ -49,7 +49,6 @@ class Inference(object):
         # Initialize configurations
         self.submodel = config.submodel
         self.inference_model = config.inference_model
-        self.sample_num = config.sample_num
         self.disable_correction = config.disable_correction
 
         # Data loader.
@@ -90,6 +89,8 @@ class Inference(object):
                                       bond_encoder=self.bond_encoder,
                                       bond_decoder=self.bond_decoder)
 
+        self.sample_num = config.sample_num if config.sample_num < self.inf_dataset.num_smiles else self.inf_dataset.num_smiles
+    
         self.inf_loader = DataLoader(self.inf_dataset,
                                  shuffle=True,
                                  batch_size=self.inf_batch_size,