Merge pull request #165 from MannLabs/change-device-charge-model

Allow change of device for charge models through initialization

peptdeep/model/charge.py

@@ -86,12 +86,21 @@ class ChargeModelForModAASeq(
 ):
     """
     ModelInterface for charge prediction for modified peptides
+
+     Parameters
+    ----------
+    min_charge : int, optional
+        Minimum charge to predict, by default 1
+    max_charge : int, optional
+        Maximum charge to predict, by default 6
+    device : str, optional
+        Device to use for training and prediction, by default "gpu"
     """
-    def __init__(self, min_charge:int=1, max_charge:int=6):
+    def __init__(self, min_charge:int=1, max_charge:int=6, device:str="gpu"):
         super().__init__(
             num_target_values=max_charge-min_charge+1,
             model_class=Model_for_Generic_ModAASeq_BinaryClassification_Transformer,
-            nlayers=4, hidden_dim=128, dropout=0.1
+            nlayers=4, hidden_dim=128, dropout=0.1, device=device
         )
 
         self.target_column_to_predict = "charge_probs"
@@ -109,12 +118,21 @@ class ChargeModelForAASeq(
 ):
     """
     ModelInterface for charge prediction for amino acid sequence
+
+    Parameters
+    ----------
+    min_charge : int, optional
+        Minimum charge to predict, by default 1
+    max_charge : int, optional
+        Maximum charge to predict, by default 6
+    device : str, optional
+        Device to use for training and prediction, by default "gpu"
     """
-    def __init__(self, min_charge:int=1, max_charge:int=6):
+    def __init__(self, min_charge:int=1, max_charge:int=6,device:str="gpu"):
         super().__init__(
             num_target_values=max_charge-min_charge+1,
             model_class=Model_for_Generic_AASeq_BinaryClassification_Transformer,
-            nlayers=4, hidden_dim=128, dropout=0.1
+            nlayers=4, hidden_dim=128, dropout=0.1, device=device
         )
 
         self.target_column_to_predict = "charge_probs"