@batch(memory=40000, gpu=1, image='<http://763104351884.dkr.ecr.us-east-1.amazonaws.com/tensorflow-training:2.14.1-gpu-py310-cu118-ubuntu20.04-ec2|763104351884.dkr.ecr.us-east-1.amazonaws.com/tensorflow-training:2.14.1-gpu-py310-cu118-ubuntu20.04-ec2>')

@conda(libraries={'numpy': '1.23.5', 'tensorflow': '2.15'})

@step

def train_model(self):