is it possible to get embedding of last layer of network and add softmax to get a Classifier

[roodkcab]

[keyu-tian]

hi @roodkcab, maybe you can add a classifier linear layer in VAR byself.head2 = nn.Linear(self.C, classification_classes) and replace self.head(self.head_nm(h.float(), cond_BD).float()).float() with self.head2(self.head_nm(h.float(), cond_BD).float().mean(dim=1)).float() in models/var.py to get your classification logits.

[roodkcab]

great! I'll have a look. really impressive work! thanks again.