Fix drive train setup