added fixes for handling multiple shape warmup#13

JRosenkranz wants to merge 2 commits intomainfoundation-model-stack/aiu-fms-testing-utils:mainfrom

multiple_shape_warmupfoundation-model-stack/aiu-fms-testing-utils:multiple_shape_warmupCopy head branch name to clipboard