In a Jenkins declarative pipeline, how can I set the agent label for dynamically created parallel stages

I have a function that is creating dynamic stages to be run in parallel, the code looks similar to this...

def stagesMap = [:]
for (int i = 1; i < 3; i++) {
    def stageName = "Stage ${i}"
    stagesMap[testName] = {
        stage("${stageName}") {
            agent { label 'SomeAgent' }
            script {
                someFunctionCall()
            }
        }
    }
}
parallel stagesMap

When I run this pipeline all the stages run on the master agents, instead of running on the agent specified in the dynamic stage.

Is there a way to specify the agent when creating stages dynamically?

The solution was to use a node, instead of a stage.

def stagesMap = [:]
for (int i = 1; i < 3; i++) {
    stagesMap[testName] = {
        node('nodeNameOrLabel') { //use node here instead of stage.
            script {
                someFunctionCall()
            }
        }
    }
}
parallel stagesMap