Upload Flowformer

model_flowformer.py

@@ -96,7 +96,7 @@ class Flowformer(PreTrainedModel):
         self.dec = nn.Sequential(*dec_layers)
 
     def markers(self):
-        return self._pretrained_markers
+        return self._markers
 
     def forward(self, tensor: torch.Tensor, labels: torch.Tensor=None, markers: list=None):
         r"""
@@ -110,7 +110,7 @@ class Flowformer(PreTrainedModel):
         """
         B, L, M = tensor.shape
         if markers is not None:
-            assert len(markers) == M, "Number of markers in x and markers must be identical"
+            assert len(markers) == M, "last dimension of input must be equal to number of markers"
 
             zeros = torch.zeros((B, L, len(self.markers())), device=tensor.device)
             valid_markers = [m for m in markers if m in set(self.markers()).intersection(markers)]