fix: get_device should gracefully handle models without a device attr

chanind · Aug 4, 2024 · f0e596c · f0e596c
1 parent fd7b86c
commit f0e596c
Showing 1 changed file with 1 addition and 1 deletion.
diff --git 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