From 5fb26819c4232e399810c36ff25601c7d3f3b02a Mon Sep 17 00:00:00 2001 From: David Chanin Date: Sun, 4 Aug 2024 08:15:17 -0600 Subject: [PATCH] fix: get_device should gracefully handle models without a device attr (#7) --- 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