Fix type of 'is_token_based' parameter

Author: daniel-unyi-42

src/segger/models/segger_model.py

@@ -31,7 +31,7 @@ def forward(self, x_dict, edge_index_dict):
 class Segger(nn.Module):
     def __init__(
         self,
-        is_token_based: int,
+        is_token_based: bool,
         num_node_features: dict[str, int],
         init_emb: int = 16,
         hidden_channels: int = 32,
@@ -43,7 +43,7 @@ def __init__(
         Initializes the Segger model.
 
         Args:
-            is_token_based (int)   : Whether the model is using token-based embeddings or scRNAseq embeddings.
+            is_token_based (bool)  : Whether the model is using token-based embeddings or scRNAseq embeddings.
             num_node_features (dict[str, int]): Number of node features for each node type.
             init_emb (int)         : Initial embedding size for both 'tx' and boundary (non-token) nodes.
             hidden_channels (int)  : Number of hidden channels.

src/segger/training/train.py

@@ -61,7 +61,7 @@ def __init__(self, learning_rate: float = 1e-3, **kwargs):
 
     def from_new(
         self,
-        is_token_based: int,
+        is_token_based: bool,
         num_node_features: dict[str, int],
         init_emb: int,
         hidden_channels: int,
@@ -75,7 +75,7 @@ def from_new(
 
         Parameters
         ----------
-        is_token_based : int
+        is_token_based : bool
             Whether the model is using token-based embeddings or scRNAseq embeddings.
         num_node_features : dict[str, int]
             Number of node features for each node type.