Makint it consistent with master so I don't have to rebase or anything similar. Also fixing bugs related to the other forms of training where I did not fix the restore weights functionality