diff --git a/dacapo/experiments/tasks/hot_distance_task.py b/dacapo/experiments/tasks/hot_distance_task.py
index 24cb2367..382eaf8b 100644
--- 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
index 2538150e..4a18a315 100644
--- 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)