Rewrite in DistributedDataParallel and test on Perlmutter
Rewrite in DistributedDataParallel and test on Perlmutter