diff --git a/cellfinder/napari/detect/detect.py b/cellfinder/napari/detect/detect.py
index 1289f5ac..66bbbb22 100644
--- a/cellfinder/napari/detect/detect.py
+++ b/cellfinder/napari/detect/detect.py
@@ -256,6 +256,7 @@ def widget(
         use_pre_trained_weights: bool,
         weights_only: bool,
         trained_model: Optional[Path],
+        model_weights,
         batch_size: int,
         misc_options,
         start_plane: int,
@@ -375,8 +376,6 @@ def widget(
             max_cluster_size,
         )
 
-        model_weights = None
-
         if weights_only:
             model_weights = trained_model
             trained_model = None
diff --git a/cellfinder/napari/detect/detect_containers.py b/cellfinder/napari/detect/detect_containers.py
index 988c9918..df18f085 100644
--- a/cellfinder/napari/detect/detect_containers.py
+++ b/cellfinder/napari/detect/detect_containers.py
@@ -131,7 +131,7 @@ def widget_representation(cls) -> dict:
             use_pre_trained_weights=dict(
                 value=cls.defaults()["use_pre_trained_weights"]
             ),
-            # weights_only=dict(value=cls.defaults()["weights_only"]),
+            weights_only=dict(value=cls.defaults()["weights_only"]),
             trained_model=dict(value=cls.defaults()["trained_model"]),
             skip_classification=dict(
                 value=cls.defaults()["skip_classification"]