# STN training crash in Apple M1

**URL:** <https://discuss.pytorch.org/t/stn-training-crash-in-apple-m1/160702>\
**Category:** Mac OS X\
**Created:** [September 4, 2022, 9:43am UTC](https://discuss.pytorch.org/t/stn-training-crash-in-apple-m1/160702 "2022-09-04T09:43:45Z")\
**Posts on this page:** 1\
**Page:** 1

<div class="post-metadata">

**Author:** ![Jack75](https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/j/3bc359/32.png) [@Jack75](https://discuss.pytorch.org/u/Jack75)\
**Post date:** [September 4, 2022, 9:43am UTC](https://discuss.pytorch.org/t/stn-training-crash-in-apple-m1/160702/1 "2022-09-04T09:43:45Z")

</div>

I am trying to speed up STN sample which is in

> **[Spatial Transformer Networks Tutorial — PyTorch Tutorials 2.9.0+cu128...](https://docs.pytorch.org/tutorials/intermediate/spatial_transformer_tutorial.html)**

I have install nighty version with conda  
$ python

Python 3.9.12 (main, Apr 5 2022, 01:52:34)

[Clang 12.0.0] :: Anaconda, Inc. on darwin

Type “help”, “copyright”, “credits” or “license” for more information.

> > > import torch

> > > import platform

> > > print(platform.platform())

macOS-12.3.1-arm64-arm-64bit

> > > torch. **version**

‘1.13.0.dev20220903’

> > > torch.device(‘mps’)

device(type=‘mps’)

> > >

I modified the code  
#device = torch.device(“cuda” if torch.cuda.is\_available() else “cpu”)

# speed up for Apple M1

device = torch.device(‘mps’)

however, the kernel is crash  
-:27:11: error: invalid input tensor shapes, indices shape and updates shape must be equal

-:27:11: note: see current operation: %25 = “mps.scatter\_along\_axis”(%23, %arg0, %24, %1) {mode = 6 : i32} : (tensor\<150528xf32\>, tensor\<28xf32\>, tensor\<50176xi32\>, tensor) → tensor\<150528xf32\>

/AppleInternal/Library/BuildRoots/560148d7-a559-11ec-8c96-4add460b61a6/Library/Caches/com.apple.xbs/Sources/MetalPerformanceShadersGraph/mpsgraph/MetalPerformanceShadersGraph/Core/Files/MPSGraphExecutable.mm:1267: failed assertion `Error: MLIR pass manager failed’