Yes this will work as long as the output backing buffer is the correct size corresponding to the size of the input.

One note is that being able to avoid data copies during inference for models with flexible shapes will vary depending on circumstances. You can use the Core ML Instrument and look in the Data lane to see if data copies are occurring.