How to use potential energy surface represented by a neural network for molecular dynamics simulation?

I trained a neural network that represents the potential energy surface of a periodic system consists of $$50$$ atoms.

The neural network takes the structure proposed by Behler and Parrinello in 2007. All $$50$$ atoms consist of one training example (structure) are within the supercell. To calculate the energy (the forces are calculated in a similar way) given a structure, the Cartesian coordinate of each atom in the system should be converted to a sub-fingerprint consists of two types of descriptors, $$G_2$$ and $$G_4$$ in the first place. Given the neighbors of an atom, $$G_2$$ is a function of atomic spacings between the center atom and the neighbors, $$G_4$$ is a function of both atomic spacings and atomic angles. The calculated fingerprint represents the whole system will be used as an input to the neural network.

After the neural network is trained, I packed it to be a readable calculator by ase, so I can utilize the MD module of ase to do molecular dynamics simulation using NVE ensemble.

However, one severe problem occurred when I analyzed the MD result, the problem is atoms will move out of the supercell in the process of MD simulation.

As I explained, fingerprints of atoms in the system should be calculated in order to calculate the energy and forces. Before calculating the fingerprint of an atom, atoms moved out of the supercell will be mirrored back by periodic boundary condition to find neighbors within a cutoff. However, actual positions of atoms will be used as atomic spacing for calculating $$G_2$$ and $$G_4$$.

I implemented two types of neural network calculator for MD simulation:

1. If an atom moved out of the supercell in the last MD step, it will be mirrored back to the supercell (the actual distances between atoms are changed) for calculating the energy and forces.

2. Since the neighbors of an atom are not affected whether atoms are mirrored back or not, so in the process of MD simulation, the positions of atoms will not be changed manually.

By using the first calculator for MD simulation (using NVE ensemble), the temperature of the system will become ridiculously high. On the contrary, the MD result from the second calculator shows atoms are moving further apart in the process (I forgot to log the process temperature).

Which calculator is reasonable? In the training examples I used, all atoms are within the supercell. I guess the second calculator is reasonable. If I'm correct, does it mean I need to include structures where some atoms are out of the supercell? (I'm using VASP)