PyTorch supports the axis keyword arg for numpy compatibility, but it looks like there is a bug where this isn’t translating into TorchScript. Most TorchScript ops use dim in place of axis (the meaning is the same), so if you change that in your code it should work, i.e. torch.cat([a, b], axis=2) becomes torch.cat([a, b], dim=2).
1 Like