Skip to content
Snippets Groups Projects

Allow specifying device to use when training

Merged Manon Blanco requested to merge export-device-when-training into main
All threads resolved!
Files
2
@@ -101,7 +101,7 @@ class GenericTrainingManager:
def init_hardware_config(self):
# Debug mode
if self.device_params["force"] not in [None, "cuda"]:
if self.device_params["force"] not in [None, "cuda"] or not torch.cuda.is_available():
self.device_params["use_ddp"] = False
self.device_params["use_amp"] = False
Loading