From f0e596c60a83834dd67e1f3ab461bcbc743a0499 Mon Sep 17 00:00:00 2001 From: David Chanin Date: Sun, 4 Aug 2024 15:11:12 +0100 Subject: [PATCH] fix: get_device should gracefully handle models without a device attr --- linear_relational/lib/torch_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/linear_relational/lib/torch_utils.py b/linear_relational/lib/torch_utils.py index 7a9ca86..a1a4eeb 100644 --- a/linear_relational/lib/torch_utils.py +++ b/linear_relational/lib/torch_utils.py @@ -22,7 +22,7 @@ def get_device(model: nn.Module) -> torch.device: """ Returns the device on which the model is running. """ - if isinstance(model.device, torch.device): + if hasattr(model, "device") and isinstance(model.device, torch.device): return model.device return next(model.parameters()).device