fixing pretrain