Pytorch model converted to neuralnetwork crashes in swift

I struggle with deploying a big model into a swift application. I need it in the neuralnetwork format since a crucial part has to be a custom layer performing calculations on gpu. Problem occurs in the standard code though and I managed to narrow it down to ~simple code. It does not calculate anything usefull, but shows the problem.

Pytorch model is:

```
import torch

def some_fn(tenInput, tenFlow):

shapeInt = torch.tensor(tenInput.shape, device=tenFlow.device)
tenOnes = torch.ones(
[ shapeInt[0], 1, shapeInt[2], shapeInt[3] ],
dtype=tenFlow.dtype,
device=tenFlow.device
)
tenOutput = torch.cat([ tenInput, tenOnes ], 1)

# here I actually apply my custom layer, but it is not needed to trigger crashes...

tenResult = tenOutput[:, :-1, :, :]
tenMask = tenOutput[:, -1:, :, :]

tenMask = torch.lt(tenMask, 0.999).expand(tenResult.shape)
tenResult[tenMask] = 0.0

return tenResult.to(torch.float32)

class TestModel(torch.nn.Module):

def forward(self, x, flow):
y = some_fn(x, flow)
return y
```

Conversion to mlmodel is also standard:

```
import coremltools as ct

# a test pattern I use to validate outputs, not really
# relevent here, any x and y inputs are ok
w = 1024
h = 768
ch = 32

x = torch.zeros(1*ch*h*w, dtype=torch.float32)
y = torch.zeros(1*2*h*w, dtype=torch.float32)
for i in range(x.shape[0]):
x[i] = 0.1 * (i % 13)
for i in range(y.shape[0]):
y[i] = (i % 19) - 10

x = x.reshape((1,ch,h,w)).to('cuda')
y = y.reshape((1,2,h,w)).to('cuda')

# conversion:
m = TestModel().to('cuda').eval()
traced_model = torch.jit.trace(m, (x, y), check_trace=False)
mlmodel = ct.convert(
traced_model,
convert_to="neuralnetwork",
inputs=[
ct.TensorType(name="x", shape=x.shape),
ct.TensorType(name="y", shape=y.shape),
],
debug=False
)
mlmodel.save("test_model.mlmodel")
```

Then I add the model to a swift project in Xcode and load/run it with the code:

```
guard let model = try? test_model() else {
fatalError("loading failed")
}

let w = 1024 as NSNumber
let h = 768 as NSNumber
let c = 32 as NSNumber

// input pattern, the same as in pytorch
guard let x_inp = try? MLMultiArray(shape:[1,c,h,w], dataType:MLMultiArrayDataType.float32) else {
fatalError("failed on frame 0")
}
guard let y_inp = try? MLMultiArray(shape:[1,2,h,w], dataType:MLMultiArrayDataType.float32) else {
fatalError("failed on frame 1")
}
for i in 0..

Start with the reduced PyTorch model, the coremltools conversion to the neuralnetwork format, and the generated Swift model.prediction call on the stated macOS and M1 setup. Compare the scatter_nd_kernel error with the reduced reproduction; done means identifying the conversion or runtime incompatibility and documenting a reproducible fix or limitation.

Written by the indexing model from the issue text.