diff --git a/napari_cellseg3d/code_models/worker_inference.py b/napari_cellseg3d/code_models/worker_inference.py
index 7936363b..c69dfe45 100644
--- a/napari_cellseg3d/code_models/worker_inference.py
+++ b/napari_cellseg3d/code_models/worker_inference.py
@@ -859,7 +859,7 @@ def inference(self):
                 model = model_class(
                     input_img_size=[dims, dims, dims],
                     # device=self.config.device,
-                    num_classes=self.config.model_info.num_classes,
+                    # num_classes=self.config.model_info.num_classes,
                 )
                 try:
                     model = model.to(self.config.device)