Revert model changes and remove print statements outside of verbose

Author: andrewmoorman

src/segger/models/segger_model.py

@@ -38,20 +38,20 @@ def __init__(
 
         # First GATv2Conv layer
         self.conv_first = GATv2Conv((-1, -1), hidden_channels, heads=heads, add_self_loops=False)
-        self.lin_first = Linear(-1, hidden_channels * heads)
+        #self.lin_first = Linear(-1, hidden_channels * heads)
 
         # Middle GATv2Conv layers
         self.num_mid_layers = num_mid_layers
         if num_mid_layers > 0:
             self.conv_mid_layers = torch.nn.ModuleList()
-            self.lin_mid_layers = torch.nn.ModuleList()
+            #self.lin_mid_layers = torch.nn.ModuleList()
             for _ in range(num_mid_layers):
                 self.conv_mid_layers.append(GATv2Conv((-1, -1), hidden_channels, heads=heads, add_self_loops=False))
-                self.lin_mid_layers.append(Linear(-1, hidden_channels * heads))
+                #self.lin_mid_layers.append(Linear(-1, hidden_channels * heads))
 
         # Last GATv2Conv layer
         self.conv_last = GATv2Conv((-1, -1), out_channels, heads=heads, add_self_loops=False)
-        self.lin_last = Linear(-1, out_channels * heads)
+        #self.lin_last = Linear(-1, out_channels * heads)
 
     def forward(self, x: Tensor, edge_index: Tensor) -> Tensor:
         """
@@ -70,19 +70,19 @@ def forward(self, x: Tensor, edge_index: Tensor) -> Tensor:
         x = self.tx_embedding(((x.sum(1) * is_one_dim).int())) * is_one_dim + self.lin0(x.float()) * (1 - is_one_dim)
         # First layer
         x = x.relu()
-        x = self.conv_first(x, edge_index)  + self.lin_first(x)
+        x = self.conv_first(x, edge_index) # + self.lin_first(x)
         x = x.relu()
 
         # Middle layers
         if self.num_mid_layers > 0:
             for i in range(self.num_mid_layers):
                 conv_mid = self.conv_mid_layers[i]
-                lin_mid  = self.lin_mid_layers[i]
-                x = conv_mid(x, edge_index)  + lin_mid(x)
+                #lin_mid  = self.lin_mid_layers[i]
+                x = conv_mid(x, edge_index) # + lin_mid(x)
                 x = x.relu()
 
         # Last layer
-        x = self.conv_last(x, edge_index)  + self.lin_last(x)
+        x = self.conv_last(x, edge_index) # + self.lin_last(x)
 
         return x
 

src/segger/prediction/predict_parquet.py

@@ -304,7 +304,6 @@ def _get_id():
         """Generate a random Xenium-style ID."""
         return "".join(np.random.choice(list("abcdefghijklmnopqrstuvwxyz"), 8)) + "-nx"
 
-    print(gpu_id)
     with cp.cuda.Device(gpu_id):
         # Move the batch to the specified GPU
         batch = batch.to(f"cuda:{gpu_id}")