Fix bug with class_label param in UVit class

Author: Alejandro Monroy Muñoz

models/early_exit.py

@@ -280,7 +280,7 @@ def forward(self, x, timesteps, y=None):
         )
         time_token = time_token.unsqueeze(dim=1)
         x = torch.cat((time_token, x), dim=1)
-        if y is not None:
+        if y is not None and self.uvit.label_emb is not None:
             label_emb = self.uvit.label_emb(y)
             label_emb = label_emb.unsqueeze(dim=1)
             x = torch.cat((label_emb, x), dim=1)

models/uvit.py

@@ -275,6 +275,7 @@ def __init__(
             self.label_emb = nn.Embedding(self.num_classes, embed_dim)
             self.extras = 2
         else:
+            self.label_emb = None
             self.extras = 1
 
         self.pos_embed = nn.Parameter(