我有这样的分类器：

input_dim = 25088
h1_dim = 4096
h2_dim = 2048
h3_dim = 1024
h4_dim = 512
output_dim = len(cat_to_name) # 102
drop_prob = 0.2

model.classifier = nn.Sequential(nn.Linear(input_dim, h1_dim),
                                 nn.ReLU(),
                                 nn.Dropout(drop_prob),
                                 nn.Linear(h1_dim, h2_dim),
                                 nn.ReLU(),
                                 nn.Dropout(drop_prob),
                                 nn.Linear(h2_dim, h3_dim),
                                 nn.ReLU(),
                                 nn.Dropout(drop_prob),
                                 nn.Linear(h3_dim, h4_dim),
                                 nn.ReLU(),
                                 nn.Dropout(drop_prob),
                                 nn.Linear(h4_dim, output_dim),
                                 )

我用CrossEntropyLoss作为标准。在验证和测试中，我如何添加Softmax？这是验证循环：

model.eval()
            with torch.no_grad():
                for images, labels in valid_loader:
                    images, labels = images.to(device), labels.to(device)
                    images.requires_grad = True

                    logits = model.forward(images)
                    batch_loss = criterion(logits, labels)
                    valid_loss += batch_loss.item()

                    ps = torch.exp(logits)
                    top_p, top_class = ps.topk(1, dim=1)
                    equals = top_class == labels.view(*top_class.shape)