fix: vit network did not conform to expected state dict

IML-DKFZ · Aug 9, 2023 · 3500f91 · 3500f91
1 parent 7545b32
commit 3500f91
Showing 1 changed file with 12 additions and 2 deletions.
diff --git a/fd_shifts/models/networks/vit.py b/fd_shifts/models/networks/vit.py
@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import timm
 import torch
 import torch.nn as nn
@@ -10,8 +12,16 @@ class ViT(Network):
     def __init__(self, cf: configs.Config):
         super().__init__()
 
-        self.encoder = Encoder(cf)
-        self.classifier = Classifier(self.encoder.model.head)
+        self._encoder = Encoder(cf)
+        self._classifier = Classifier(self.encoder.model.head)
+
+    @property
+    def encoder(self) -> Encoder:
+        return self._encoder
+
+    @property
+    def classifier(self) -> Classifier:
+        return self._classifier
 
     def forward(self, x):
         out = self.encoder(x)