training function now takes a learning rate
removed faulty code from train mehtod