diff --git a/models/__init__.py b/models/__init__.py index 8d40c77d8f59b1fc8ff987a716a0f1ea2182fd4f..04b11b2e094282f8fe91c0291382788bd8fe9dec 100755 --- a/models/__init__.py +++ b/models/__init__.py @@ -75,9 +75,5 @@ def create_model(pretrained, dataset, arch, parallel=True, device_ids=None): elif parallel: model = torch.nn.DataParallel(model, device_ids=device_ids) - # explicitly add a softmax layer, because it is useful when exporting to ONNX - model.original_forward = model.forward - softmax = torch.nn.Softmax(dim=1) - model.forward = lambda input: softmax(model.original_forward(input)) model.cuda() return model