When training rectified_flow_2, the code rectified_flow_2.net = copy.deepcopy(rectified_flow_1) is wrong and useless. In fact, I think you should use rectified_flow_2.model.load_state_dict(rectified_flow_1.model.state_dict()) to load pretrained parameters of rectified_flow_1.