Prediction of fluid flow in porous media possess substantial engineering significance. Physics-Informed Neural Networks (PINNs) is a promising method to solve Partial Differential Equations (PDEs) by leveraging the universal approximation ability of Neural Network (NN). However, current PINN-based methods are rarely used to simulate strongly heterogeneous problems, where the combination of accuracy and efficiency is challenging. In this work, we introduce a novel training strategy to enhance the performance of the Enriched Physics-Informed Neural Network (E-PINN) for two-phase flow simulation. Specifically, the PDEs residual is constructed by Finite Volume Method (FVM), boundary conditions are incorporated into the training process as hard constraints, permeability fields of various degrees of heterogeneity are generated by Stanford Geostatistical Modeling Software (SGeMS), the NN is trained using a physics-driven approach without labeled data. Innovatively, a pruning strategy is designed for model training to enhance model’s convergence ability in heterogeneous scenarios. We simulate multiple test cases to verify the superiority of the pruned E-PINN, the new method is proven to achieve sufficient accuracy while significantly improving training speed.