Fixed an error. Added call to Create in train_task