bugfix: kernel size

janelia-cellmap · Nov 19, 2024 · abfcd02 · abfcd02
1 parent c52ee0f
commit abfcd02
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/dacapo/experiments/tasks/hot_distance_task.py b/dacapo/experiments/tasks/hot_distance_task.py
@@ -47,6 +47,7 @@ def __init__(self, task_config):
             channels=task_config.channels,
             scale_factor=task_config.scale_factor,
             mask_distances=task_config.mask_distances,
+            kernel_size=task_config.kernel_size,
         )
         self.loss = HotDistanceLoss()
         self.post_processor = ThresholdPostProcessor()

diff --git a/dacapo/experiments/tasks/predictors/hot_distance_predictor.py b/dacapo/experiments/tasks/predictors/hot_distance_predictor.py
@@ -120,11 +120,11 @@ def create_model(self, architecture):
         """
         if architecture.dims == 2:
             head = torch.nn.Conv2d(
-                architecture.num_out_channels, self.embedding_dims, self.kernel_size=1
+                architecture.num_out_channels, self.embedding_dims, self.kernel_size
             )
         elif architecture.dims == 3:
             head = torch.nn.Conv3d(
-                architecture.num_out_channels, self.embedding_dims, self.kernel_size=1
+                architecture.num_out_channels, self.embedding_dims, self.kernel_size
             )
 
         return Model(architecture, head)