custom gpu option

README.md

@@ -46,7 +46,7 @@ pip install rdkit
 
 ```python
 from rxnmapper import RXNMapper
-rxn_mapper = RXNMapper()
+rxn_mapper = RXNMapper(gpu_id=1) # to specify which gpu to use to load the model, in case of multi gpu system.
 rxns = ['CC(C)S.CN(C)C=O.Fc1cccnc1F.O=C([O-])[O-].[K+].[K+]>>CC(C)Sc1ncccc1F', 'C1COCCO1.CC(C)(C)OC(=O)CONC(=O)NCc1cccc2ccccc12.Cl>>O=C(O)CONC(=O)NCc1cccc2ccccc12']
 results = rxn_mapper.get_attention_guided_atom_maps(rxns)
 ```
@@ -63,7 +63,7 @@ The results contain the mapped reactions and confidence scores:
 To account for batching and error handling automatically, you can use `BatchedMapper` instead:
 ```python
 from rxnmapper import BatchedMapper
-rxn_mapper = BatchedMapper(batch_size=32)
+rxn_mapper = BatchedMapper(batch_size=32, gpu_id=1) # to specify which gpu to use to load the model, in case of multi gpu system.
 rxns = ['CC[O-]~[Na+].BrCC>>CCOCC', 'invalid>>reaction']
 
 # The following calls work with input of arbitrary size. Also, they do not raise 

rxnmapper/batched_mapper.py

@@ -29,6 +29,7 @@ def __init__(
         model_type: str = "albert",
         canonicalize: bool = False,
         placeholder_for_invalid: str = ">>",
+        gpu_id: int = 0
     ):
         """
         Args:
@@ -56,6 +57,7 @@ def __init__(
                 model_type=model_type,
                 attention_multiplier=attention_multiplier,
             ),
+            gpu_id=gpu_id
         )
         self.batch_size = batch_size
         self.canonicalize = canonicalize

rxnmapper/core.py

@@ -38,6 +38,7 @@ def __init__(
         self,
         config: Optional[Dict[str, Any]] = None,
         logger: Optional[logging.Logger] = None,
+        gpu_id: int = 0
     ):
         """
         RXNMapper constructor.
@@ -70,7 +71,7 @@ def __init__(
 
         self.logger = logger if logger else _logger
         self.model, self.tokenizer = self._load_model_and_tokenizer()
-        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.device = torch.device("cuda:{}".format(gpu_id) if torch.cuda.is_available() else "cpu")
         self.model.to(self.device)
 
     def _load_model_and_tokenizer(self) -> Tuple: