Clean-Up and Modifications towards the latest merge !22

• Update the imports for running as module
• update the usage of Dataset and Dataloader within a test
• remove unnecessary tests
• remove other dead code