maxvit fix

mwalmsley · Feb 6, 2024 · 3095b5a · 3095b5a
1 parent adcce7d
commit 3095b5a
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/zoobot/pytorch/training/finetune.py b/zoobot/pytorch/training/finetune.py
@@ -175,7 +175,7 @@ def configure_optimizers(self):
                 self.encoder.layer4
             ]
         elif isinstance(self.encoder, timm.models.MaxxVit):
-            blocks_to_tune = self.encoder.stem + [stage for stage in self.encoder.stages]
+            blocks_to_tune = [self.encoder.stem] + [stage for stage in self.encoder.stages]
         else:
             raise ValueError(f'Encoder architecture not automatically recognised: {type(self.encoder)}')