docs

mwalmsley · Feb 27, 2024 · 772660d · 772660d
1 parent b2f8e39
commit 772660d
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/zoobot/pytorch/training/finetune.py b/zoobot/pytorch/training/finetune.py
@@ -177,7 +177,7 @@ def configure_optimizers(self):
             ]
         elif isinstance(self.encoder, timm.models.MaxxVit):
             blocks_to_tune = [self.encoder.stem] + [stage for stage in self.encoder.stages]
-        elif isinstance(self.encoder, timm.models.ConvNeXt):  # stem + 3 blocks, for all sizes
+        elif isinstance(self.encoder, timm.models.ConvNeXt):  # stem + 4 blocks, for all sizes
             # https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/convnext.py#L264
             blocks_to_tune = [self.encoder.stem] + [stage for stage in self.encoder.stages]
         else: