16 lines
645 B
Python
16 lines
645 B
Python
from pytorch_lightning.accelerators.accelerator import Accelerator
|
|
from pytorch_lightning.plugins.precision import MixedPrecisionPlugin
|
|
from pytorch_lightning.utilities.exceptions import MisconfigurationException
|
|
|
|
|
|
class CPUAccelerator(Accelerator):
|
|
|
|
def setup(self, trainer, model):
|
|
if isinstance(self.precision_plugin, MixedPrecisionPlugin):
|
|
MisconfigurationException("amp + cpu is not supported. Please use a GPU option")
|
|
|
|
if "cpu" not in str(self.root_device):
|
|
raise MisconfigurationException(f"Device should be CPU, got {self.root_device} instead")
|
|
|
|
return super().setup(trainer, model)
|