diff --git a/LION/optimizers/LIONsolver.py b/LION/optimizers/LIONsolver.py index 22fa8a54..5d6832f1 100644 --- a/LION/optimizers/LIONsolver.py +++ b/LION/optimizers/LIONsolver.py @@ -236,7 +236,9 @@ def check_training_ready(self, error=True, autofill=True, verbose=True): error=False, autofill=autofill, verbose=verbose, - default=torch.device(torch.cuda.current_device()), + default=torch.device( + torch.cuda.current_device() if torch.cuda.is_available() else "cpu" + ), ) # Test 2: is the model set? if not, raise error or warn