https://github.com/vrvrv/PyTorch-SBM/blob/main/src/models/ncsn.py