Reviewed DataModule design+docs+types

In this MR, I try to review the use of the data module and go further in simplifications for the whole code:

• Slightly streamlined the datamodule approach
• Added documentation
• Added type annotations
• Added some TODOs for further discussion

Closes #7 (closed).

TODO:

• (@dcarron) Implement/document/type hint normaliser (ensure it is used correctly through all the models; integrate on the train script) and that train CLI works!
• (@andre.anjos) Deal with import reweight_BCEWithLogitsLoss

src/ptbench/utils/accelerator.py

@@ -18,8 +18,10 @@ class AcceleratorProcessor:
     """
 
     def __init__(self, name):
-        # Note: "auto" is a valid accelerator in lightning, but there doesn't seem to be a way to check which accelerator it will actually use so we don't take it into account for now.
-        self.torch_to_lightning = {"cpu": "cpu", "cuda": "gpu"}
+        # Note: "auto" is a valid accelerator in lightning, but there doesn't
+        # seem to be a way to check which accelerator it will actually use so
+        # we don't take it into account for now.
+        self.torch_to_lightning = {"cpu": "cpu", "cuda": "gpu", "mps": "mps"}
 
         self.lightning_to_torch = {
             v: k for k, v in self.torch_to_lightning.items()
@@ -57,6 +59,8 @@ class AcceleratorProcessor:
                         "Environment variable 'CUDA_VISIBLE_DEVICES' is not set."
                         "Please set 'CUDA_VISIBLE_DEVICES' of specify a device to use, e.g. cuda:0"
                     )
+        elif self.accelerator == "mps":
+            self.device = 1
         else:
             # No need to check the CUDA_VISIBLE_DEVICES environment variable if cpu
             pass