Updatable PyTorch model

Hello, I am trying to convert an off the shelf PyTorch model to be updatable. The model I am using with CoreML doesn't have a softmax layer at the end because that's handled by the loss function. Setting the cross entropy loss using the coremltools build however seems to require a softmax output. I tried adding one like this:
```
# Load a pre-trained version of MobileNetV2 model.
torch_model = torchvision.models.mobilenet_v2(weights=torchvision.models.MobileNet_V2_Weights.DEFAULT)

torch_model = torch.nn.Sequential(
torch_model,
torch.nn.Softmax(dim=1)
)

torch_model.eval()
```

but it get's converted as a softmax_nd instead of a normal softmax which the updatable script doesn't seem to like. Not sure how to make it the PyTorch layer a normal softmax?

What would be the appropriate way to set this model to be updatable?

Thank you.

No repository file, test, or entry point is named. Start by reproducing the MobileNetV2 conversion with the added torch.nn.Softmax(dim=1) and inspect how coremltools represents the resulting softmax_nd operation; done means establishing a supported path for an updatable model or documenting the limitation.

Written by the indexing model from the issue text.