Made changes to driver class