Add wrapper module and merge it with training module
Remove redundant training module