Updates to the data module and models

src/ptbench/data/datamodule.py

src/ptbench/engine/trainer.py

src/ptbench/models/densenet.py

src/ptbench/models/pasa.py

src/ptbench/scripts/train.py