Exploding Validation Loss

[lsaeuro]

Hello Everyone,
I am having trouble training a pytorch-lightning implementation of RandLa-Net.
I observe a validation loss of a different order of magnitude with respect to the training one :

The above image is the result of the following sanity check :

• Same data for training and Validation
• Same Code in training_step and validation_step
• Batch Norm left with track_running_stats = True ( with track_running_stats = False I obtain the same loss in val and train )

The tensors start changing order of magnitude from the first batch norm layer.
What can I do to solve this issue and what can be useful to post here for help?
Here a Dumb code not ready for running https://github.com/lsaeuro/RandLa-Dump