Upload model

modeling_vitmix.py

@@ -188,7 +188,7 @@ class ViTMixModel(PreTrainedModel):
                 mlp_dim = config.mlp_dim,
                 num_experts = config.num_experts
             )
-    def forward(self,tensor):
+    def forward(self,tensor, labels = None):
         logits = self.model(tensor)
         if labels is not None:
             loss = torch.nn.cross_entropy(logits, labels)