The optimizer part works out of the box, yes. But my training is usually consuming a yaml file with parameters and not via argparse. This is the part I had to adjust