Skip to content
Snippets Groups Projects

Making use of LightningDataModule and simplification of data loading

Merged Making use of LightningDataModule and simplification of data loading
add-datamodule into main
Merged Daniel CARRON requested to merge add-datamodule into main
Viewing commit a0f264f0
Show latest version
1 file
+ 6
− 2
Compare changes
  • Side-by-side
  • Inline
src/ptbench/utils/accelerator.py
+ 6
− 2
@@ -18,8 +18,10 @@ class AcceleratorProcessor:
@@ -18,8 +18,10 @@ class AcceleratorProcessor:
"""
"""
def __init__(self, name):
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.
# Note: "auto" is a valid accelerator in lightning, but there doesn't
self.torch_to_lightning = {"cpu": "cpu", "cuda": "gpu"}
# 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 = {
self.lightning_to_torch = {
v: k for k, v in self.torch_to_lightning.items()
v: k for k, v in self.torch_to_lightning.items()
@@ -57,6 +59,8 @@ class AcceleratorProcessor:
@@ -57,6 +59,8 @@ class AcceleratorProcessor:
"Environment variable 'CUDA_VISIBLE_DEVICES' is not set."
"Environment variable 'CUDA_VISIBLE_DEVICES' is not set."
"Please set 'CUDA_VISIBLE_DEVICES' of specify a device to use, e.g. cuda:0"
"Please set 'CUDA_VISIBLE_DEVICES' of specify a device to use, e.g. cuda:0"
)
)
 
elif self.accelerator == "mps":
 
self.device = 1
else:
else:
# No need to check the CUDA_VISIBLE_DEVICES environment variable if cpu
# No need to check the CUDA_VISIBLE_DEVICES environment variable if cpu
pass
pass